Streamline traits and pre-/post-save hooks

README.rst

@@ -238,7 +238,7 @@ This is how you could define a factory for `Django`_'s built-in ``User`` model.
             instance.is_active = True
 
         @pre_save
-        def __set_password(instance):
+        def _set_password(self, instance: User) -> None:
             instance.set_password("test")
 
 And this is how you could use it:

docs/factories.rst

@@ -38,6 +38,7 @@ Django example
         FileSystemStorage,
         SubFactory,
         pre_save,
+        trait,
     )
 
     from article.models import Article
@@ -64,11 +65,16 @@ Django example
             model = User
             get_or_create = ("username",)
 
-        @staticmethod
         @pre_save
-        def __set_password(instance):
+        def _set_password(self, instance):
             instance.set_password("test")
 
+        @trait
+        def is_admin_user(self, instance: User) -> None:
+            instance.is_superuser = True
+            instance.is_staff = True
+            instance.is_active = True
+
     class ArticleFactory(DjangoModelFactory):
 
         title = FACTORY.sentence()
@@ -96,6 +102,9 @@ Django example
     # Create one article with authors username set to admin.
     article = ArticleFactory(author__username="admin")
 
+    # Using trait
+    user = UserFactory(is_admin_user=True)
+
 Pydantic example
 ----------------
 **article/models.py**
@@ -167,6 +176,16 @@ Pydantic example
         class Meta:
             model = User
 
+        @trait
+        def is_admin_user(self, instance: User) -> None:
+            instance.is_superuser = True
+            instance.is_staff = True
+            instance.is_active = True
+
+        @pre_save
+        def _set_password(self, instance):
+            instance.set_password("test")
+
     class ArticleFactory(ModelFactory):
         id = FACTORY.pyint()
         title = FACTORY.sentence()
@@ -259,6 +278,16 @@ TortoiseORM example
             model = User
             get_or_create = ("username",)
 
+        @trait
+        def is_admin_user(self, instance: User) -> None:
+            instance.is_superuser = True
+            instance.is_staff = True
+            instance.is_active = True
+
+        @pre_save
+        def _set_password(self, instance):
+            instance.set_password("test")
+
     class ArticleFactory(TortoiseModelFactory):
         """Article factory."""
 
@@ -349,6 +378,16 @@ Dataclasses example
         class Meta:
             model = User
 
+        @trait
+        def is_admin_user(self, instance: User) -> None:
+            instance.is_superuser = True
+            instance.is_staff = True
+            instance.is_active = True
+
+        @pre_save
+        def _set_password(self, instance):
+            instance.set_password("test")
+
     class ArticleFactory(ModelFactory):
         id = FACTORY.pyint()
         title = FACTORY.sentence()
@@ -362,3 +401,5 @@ Dataclasses example
 
         class Meta:
             model = Article
+
+*Used just like in previous example.*

examples/dataclasses/article/factories.py

@@ -45,14 +45,12 @@ def is_admin_user(self, instance: User) -> None:
         instance.is_staff = True
         instance.is_active = True
 
-    @staticmethod
     @pre_save
-    def __pre_save_method(instance):
+    def _pre_save_method(self, instance):
         instance.pre_save_called = True
 
-    @staticmethod
     @post_save
-    def __post_save_method(instance):
+    def _post_save_method(self, instance):
         instance.post_save_called = True
 
 
@@ -70,12 +68,10 @@ class ArticleFactory(ModelFactory):
     class Meta:
         model = Article
 
-    @staticmethod
     @pre_save
-    def __pre_save_method(instance):
+    def _pre_save_method(self, instance):
         instance.pre_save_called = True
 
-    @staticmethod
     @post_save
-    def __post_save_method(instance):
+    def _post_save_method(self, instance):
         instance.post_save_called = True

examples/django/article/factories.py

@@ -51,19 +51,16 @@ def is_admin_user(self, instance: User) -> None:
         instance.is_staff = True
         instance.is_active = True
 
-    @staticmethod
     @pre_save
-    def __set_password(instance):
+    def _set_password(self, instance):
         instance.set_password("test")
 
-    @staticmethod
     @pre_save
-    def __pre_save_method(instance):
+    def _pre_save_method(self, instance):
         instance.pre_save_called = True
 
-    @staticmethod
     @post_save
-    def __post_save_method(instance):
+    def _post_save_method(self, instance):
         instance.post_save_called = True
 
 
@@ -88,12 +85,10 @@ class ArticleFactory(DjangoModelFactory):
     class Meta:
         model = Article
 
-    @staticmethod
     @pre_save
-    def __pre_save_method(instance):
+    def _pre_save_method(self, instance):
         instance.pre_save_called = True
 
-    @staticmethod
     @post_save
-    def __post_save_method(instance):
+    def _post_save_method(self, instance):
         instance.post_save_called = True

examples/pydantic/article/factories.py

@@ -45,14 +45,12 @@ def is_admin_user(self, instance: User) -> None:
         instance.is_staff = True
         instance.is_active = True
 
-    @staticmethod
     @pre_save
-    def __pre_save_method(instance):
+    def _pre_save_method(self, instance):
         instance.pre_save_called = True
 
-    @staticmethod
     @post_save
-    def __post_save_method(instance):
+    def _post_save_method(self, instance):
         instance.post_save_called = True
 
 
@@ -70,12 +68,10 @@ class ArticleFactory(ModelFactory):
     class Meta:
         model = Article
 
-    @staticmethod
     @pre_save
-    def __pre_save_method(instance):
+    def _pre_save_method(self, instance):
         instance.pre_save_called = True
 
-    @staticmethod
     @post_save
-    def __post_save_method(instance):
+    def _post_save_method(self, instance):
         instance.post_save_called = True

examples/tortoise/article/factories.py

@@ -47,14 +47,12 @@ def is_admin_user(self, instance: User) -> None:
         instance.is_staff = True
         instance.is_active = True
 
-    @staticmethod
     @pre_save
-    def __pre_save_method(instance):
+    def _pre_save_method(self, instance):
         instance.pre_save_called = True
 
-    @staticmethod
     @post_save
-    def __post_save_method(instance):
+    def _post_save_method(self, instance):
         instance.post_save_called = True
 
 
@@ -73,12 +71,10 @@ class ArticleFactory(TortoiseModelFactory):
     class Meta:
         model = Article
 
-    @staticmethod
     @pre_save
-    def __pre_save_method(instance):
+    def _pre_save_method(self, instance):
         instance.pre_save_called = True
 
-    @staticmethod
     @post_save
-    def __post_save_method(instance):
+    def _post_save_method(self, instance):
         instance.post_save_called = True

fake.py

@@ -1643,7 +1643,7 @@ def __init_subclass__(cls, **kwargs):
     @classmethod
     def _run_hooks(cls, hooks, instance):
         for method in hooks:
-            getattr(cls, method)(instance)
+            getattr(cls, method)(cls, instance)
 
     @classmethod
     def _apply_traits(cls, instance, **kwargs) -> None:
@@ -2516,14 +2516,12 @@ def is_admin_user(self, instance: User) -> None:
                 instance.is_staff = True
                 instance.is_active = True
 
-            @staticmethod
             @pre_save
-            def __pre_save_method(instance):
+            def _pre_save_method(self, instance):
                 instance.pre_save_called = True
 
-            @staticmethod
             @post_save
-            def __post_save_method(instance):
+            def _post_save_method(self, instance):
                 instance.post_save_called = True
 
         class ArticleFactory(ModelFactory):
@@ -2595,14 +2593,12 @@ def is_admin_user(self, instance: User) -> None:
                 instance.is_staff = True
                 instance.is_active = True
 
-            @staticmethod
             @pre_save
-            def __pre_save_method(instance):
+            def _pre_save_method(self, instance):
                 instance.pre_save_called = True
 
-            @staticmethod
             @post_save
-            def __post_save_method(instance):
+            def _post_save_method(self, instance):
                 instance.post_save_called = True
 
         class DjangoArticleFactory(DjangoModelFactory):
@@ -2619,14 +2615,12 @@ class DjangoArticleFactory(DjangoModelFactory):
             class Meta:
                 model = Article
 
-            @staticmethod
             @pre_save
-            def __pre_save_method(instance):
+            def _pre_save_method(self, instance):
                 instance.pre_save_called = True
 
-            @staticmethod
             @post_save
-            def __post_save_method(instance):
+            def _post_save_method(self, instance):
                 instance.post_save_called = True
 
         django_article = DjangoArticleFactory(author__username="admin")
@@ -2775,14 +2769,12 @@ def is_admin_user(self, instance: TortoiseUser) -> None:
                 instance.is_staff = True
                 instance.is_active = True
 
-            @staticmethod
             @pre_save
-            def __pre_save_method(instance):
+            def _pre_save_method(self, instance):
                 instance.pre_save_called = True
 
-            @staticmethod
             @post_save
-            def __post_save_method(instance):
+            def _post_save_method(self, instance):
                 instance.post_save_called = True
 
         class TortoiseArticleFactory(TortoiseModelFactory):
@@ -2799,14 +2791,12 @@ class TortoiseArticleFactory(TortoiseModelFactory):
             class Meta:
                 model = TortoiseArticle
 
-            @staticmethod
             @pre_save
-            def __pre_save_method(instance):
+            def _pre_save_method(self, instance):
                 instance.pre_save_called = True
 
-            @staticmethod
             @post_save
-            def __post_save_method(instance):
+            def _post_save_method(self, instance):
                 instance.post_save_called = True
 
         tortoise_article = TortoiseArticleFactory(author__username="admin")