From 0c7c1768b4032ac70d0b60c8afaa70f08b171964 Mon Sep 17 00:00:00 2001 From: Artur Barseghyan Date: Wed, 6 Dec 2023 00:54:33 +0100 Subject: [PATCH] Streamline traits and pre-/post-save hooks --- README.rst | 2 +- docs/factories.rst | 45 ++++++++++++++++++++++- examples/dataclasses/article/factories.py | 12 ++---- examples/django/article/factories.py | 15 +++----- examples/pydantic/article/factories.py | 12 ++---- examples/tortoise/article/factories.py | 12 ++---- fake.py | 32 ++++++---------- 7 files changed, 72 insertions(+), 58 deletions(-) diff --git a/README.rst b/README.rst index 322961b..fbcc0dd 100644 --- a/README.rst +++ b/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: diff --git a/docs/factories.rst b/docs/factories.rst index 131a37b..08e33a5 100644 --- a/docs/factories.rst +++ b/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.* diff --git a/examples/dataclasses/article/factories.py b/examples/dataclasses/article/factories.py index 4e31a28..881b329 100644 --- a/examples/dataclasses/article/factories.py +++ b/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 diff --git a/examples/django/article/factories.py b/examples/django/article/factories.py index 060a066..d4c9c75 100644 --- a/examples/django/article/factories.py +++ b/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 diff --git a/examples/pydantic/article/factories.py b/examples/pydantic/article/factories.py index 4e31a28..881b329 100644 --- a/examples/pydantic/article/factories.py +++ b/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 diff --git a/examples/tortoise/article/factories.py b/examples/tortoise/article/factories.py index 3b4c7c8..02d4acc 100644 --- a/examples/tortoise/article/factories.py +++ b/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 diff --git a/fake.py b/fake.py index 891ebc0..2ef9920 100644 --- a/fake.py +++ b/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")