Feature/add traits

CHANGELOG.rst

@@ -16,6 +16,14 @@ are used for versioning (schema follows below):
   0.3.4 to 0.4).
 - All backwards incompatible changes are mentioned in this document.
 
+0.3.1
+-----
+2023-12-04
+
+- Improve Tortoise ORM factory.
+- Add traits.
+- Improve docmentation.
+
 0.3
 ---
 2023-12-03

Makefile

@@ -1,5 +1,5 @@
 # Update version ONLY here
-VERSION := 0.3
+VERSION := 0.3.1
 SHELL := /bin/bash
 # Makefile for project
 VENV := ~/.virtualenvs/fake.py/bin/activate

examples/dataclasses/article/factories.py

@@ -7,6 +7,7 @@
     SubFactory,
     post_save,
     pre_save,
+    trait,
 )
 
 from article.models import Article, User
@@ -38,6 +39,12 @@ class UserFactory(ModelFactory):
     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
+
     @staticmethod
     @pre_save
     def __pre_save_method(instance):

examples/django/article/factories.py

@@ -7,6 +7,7 @@
     SubFactory,
     post_save,
     pre_save,
+    trait,
 )
 
 from article.models import Article
@@ -44,6 +45,12 @@ class Meta:
         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
+
     @staticmethod
     @pre_save
     def __set_password(instance):

examples/pydantic/article/factories.py

@@ -7,6 +7,7 @@
     SubFactory,
     post_save,
     pre_save,
+    trait,
 )
 
 from article.models import Article, User
@@ -38,6 +39,12 @@ class UserFactory(ModelFactory):
     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
+
     @staticmethod
     @pre_save
     def __pre_save_method(instance):

examples/tortoise/article/factories.py

@@ -7,6 +7,7 @@
     TortoiseModelFactory,
     post_save,
     pre_save,
+    trait,
 )
 
 from article.models import Article, User
@@ -40,6 +41,12 @@ class Meta:
         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
+
     @staticmethod
     @pre_save
     def __pre_save_method(instance):

fake.py

@@ -35,7 +35,7 @@
 )
 
 __title__ = "fake.py"
-__version__ = "0.3"
+__version__ = "0.3.1"
 __author__ = "Artur Barseghyan <artur.barseghyan@gmail.com>"
 __copyright__ = "2023 Artur Barseghyan"
 __license__ = "MIT"
@@ -61,6 +61,7 @@
     "post_save",
     "pre_save",
     "run_async_in_thread",
+    "trait",
 )
 
 LOGGER = logging.getLogger(__name__)
@@ -1625,6 +1626,11 @@ def post_save(func):
     return func
 
 
+def trait(func):
+    func.is_trait = True
+    return func
+
+
 class ModelFactory:
     """ModelFactory."""
 
@@ -1654,20 +1660,42 @@ def _run_hooks(cls, hooks, instance):
         for method in hooks:
             getattr(cls, method)(instance)
 
+    @classmethod
+    def _apply_traits(cls, instance, **kwargs) -> None:
+        for name, method in cls.__dict__.items():
+            if getattr(method, "is_trait", False) and kwargs.get(name, False):
+                method(cls, instance)
+
     @classmethod
     def create(cls, **kwargs):
+        trait_keys = {
+            name
+            for name, method in cls.__dict__.items()
+            if getattr(method, "is_trait", False)
+        }
+
         model_data = {
             field: (
                 value()
                 if isinstance(value, (FactoryMethod, SubFactory))
                 else value
             )
             for field, value in cls.__dict__.items()
-            if not field.startswith("_") and not field == "Meta"
+            if (
+                not field.startswith("_")
+                and not field == "Meta"
+                and not getattr(value, "is_trait", False)
+                and not getattr(value, "is_pre_save", False)
+                and not getattr(value, "is_post_save", False)
+            )
         }
-        model_data.update(kwargs)
+        # Update model_data with non-trait kwargs
+        model_data.update(
+            {k: v for k, v in kwargs.items() if k not in trait_keys}
+        )
 
         instance = cls.Meta.model(**model_data)
+        cls._apply_traits(instance, **kwargs)
 
         pre_save_hooks = [
             method
@@ -1725,7 +1753,13 @@ def create(cls, **kwargs):
         model_data = {
             field: value
             for field, value in cls.__dict__.items()
-            if not field.startswith("_") and not field == "Meta"
+            if (
+                not field.startswith("_")
+                and not field == "Meta"
+                and not getattr(value, "is_trait", False)
+                and not getattr(value, "is_pre_save", False)
+                and not getattr(value, "is_post_save", False)
+            )
         }
 
         # Separate nested attributes and direct attributes
@@ -1743,6 +1777,7 @@ def create(cls, **kwargs):
 
         # Create a new instance if none found
         instance = cls.Meta.model(**model_data)
+        cls._apply_traits(instance, **kwargs)
 
         # Handle nested attributes
         for attr, value in nested_attrs.items():
@@ -1826,7 +1861,13 @@ async def async_filter():
         model_data = {
             field: value
             for field, value in cls.__dict__.items()
-            if not field.startswith("_") and not field == "Meta"
+            if (
+                not field.startswith("_")
+                and not field == "Meta"
+                and not getattr(value, "is_trait", False)
+                and not getattr(value, "is_pre_save", False)
+                and not getattr(value, "is_post_save", False)
+            )
         }
 
         # Separate nested attributes and direct attributes
@@ -1844,6 +1885,7 @@ async def async_filter():
 
         # Create a new instance if none found
         instance = cls.Meta.model(**model_data)
+        cls._apply_traits(instance, **kwargs)
 
         # Handle nested attributes
         for attr, value in nested_attrs.items():
@@ -2483,6 +2525,12 @@ class UserFactory(ModelFactory):
             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
+
             @staticmethod
             @pre_save
             def __pre_save_method(instance):
@@ -2528,6 +2576,14 @@ class Meta:
             hasattr(user, "post_save_called") and user.post_save_called
         )
 
+        # Testing traits
+        admin_user = UserFactory(is_admin_user=True)
+        self.assertTrue(
+            admin_user.is_staff
+            and admin_user.is_superuser
+            and admin_user.is_active
+        )
+
         # **********************************
         # ******* DjangoModelFactory *******
         # **********************************
@@ -2548,6 +2604,12 @@ class Meta:
                 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
+
             @staticmethod
             @pre_save
             def __pre_save_method(instance):
@@ -2613,6 +2675,14 @@ def __post_save_method(instance):
         self.assertEqual(len(django_articles), 5)
         self.assertIsInstance(django_articles[0], Article)
 
+        # Testing traits
+        django_admin_user = DjangoUserFactory(is_admin_user=True)
+        self.assertTrue(
+            django_admin_user.is_staff
+            and django_admin_user.is_superuser
+            and django_admin_user.is_active
+        )
+
         # **********************************
         # ******* TortoiseModelFactory *******
         # **********************************
@@ -2714,6 +2784,12 @@ class Meta:
                 model = TortoiseUser
                 get_or_create = ("username",)
 
+            @trait
+            def is_admin_user(self, instance: TortoiseUser) -> None:
+                instance.is_superuser = True
+                instance.is_staff = True
+                instance.is_active = True
+
             @staticmethod
             @pre_save
             def __pre_save_method(instance):
@@ -2779,6 +2855,14 @@ def __post_save_method(instance):
         self.assertEqual(len(tortoise_articles), 5)
         self.assertIsInstance(tortoise_articles[0], TortoiseArticle)
 
+        # Testing traits
+        tortoise_admin_user = TortoiseUserFactory(is_admin_user=True)
+        self.assertTrue(
+            tortoise_admin_user.is_staff
+            and tortoise_admin_user.is_superuser
+            and tortoise_admin_user.is_active
+        )
+
     def test_registry_integration(self) -> None:
         """Test `add`."""
         # Create a TXT file.

pyproject.toml

@@ -2,7 +2,7 @@
 name = "fake.py"
 description = "Minimalistic, standalone alternative fake data generator with no dependencies."
 readme = "README.rst"
-version = "0.3"
+version = "0.3.1"
 dependencies = []
 authors = [
     {name = "Artur Barseghyan", email = "artur.barseghyan@gmail.com"},
@@ -134,7 +134,7 @@ ignore-path = [
 ]
 
 [tool.pytest.ini_options]
-minversion = "0.3"
+minversion = "0.3.1"
 addopts = [
     "-ra",
     "-vvv",