From f148a71b779bcef184181e885b782373f5293b89 Mon Sep 17 00:00:00 2001 From: Artur Barseghyan Date: Mon, 4 Dec 2023 23:02:26 +0100 Subject: [PATCH] Feature/add traits (#12) * Add traits --- CHANGELOG.rst | 8 ++ Makefile | 2 +- examples/dataclasses/article/factories.py | 7 ++ examples/django/article/factories.py | 7 ++ examples/pydantic/article/factories.py | 7 ++ examples/tortoise/article/factories.py | 7 ++ fake.py | 94 +++++++++++++++++++++-- pyproject.toml | 4 +- 8 files changed, 128 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index fb5e675..a2a64d9 100644 --- a/CHANGELOG.rst +++ b/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 diff --git a/Makefile b/Makefile index e278640..972675a 100644 --- a/Makefile +++ b/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 diff --git a/examples/dataclasses/article/factories.py b/examples/dataclasses/article/factories.py index 23e30f1..4e31a28 100644 --- a/examples/dataclasses/article/factories.py +++ b/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): diff --git a/examples/django/article/factories.py b/examples/django/article/factories.py index 44d8043..060a066 100644 --- a/examples/django/article/factories.py +++ b/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): diff --git a/examples/pydantic/article/factories.py b/examples/pydantic/article/factories.py index 23e30f1..4e31a28 100644 --- a/examples/pydantic/article/factories.py +++ b/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): diff --git a/examples/tortoise/article/factories.py b/examples/tortoise/article/factories.py index e612b59..3b4c7c8 100644 --- a/examples/tortoise/article/factories.py +++ b/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): diff --git a/fake.py b/fake.py index e4d3609..92a9d88 100644 --- a/fake.py +++ b/fake.py @@ -35,7 +35,7 @@ ) __title__ = "fake.py" -__version__ = "0.3" +__version__ = "0.3.1" __author__ = "Artur Barseghyan " __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,8 +1660,20 @@ 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() @@ -1663,11 +1681,21 @@ def create(cls, **kwargs): 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. diff --git a/pyproject.toml b/pyproject.toml index 36893ff..5886b98 100644 --- a/pyproject.toml +++ b/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",