Repository: litestar-org/polyfactory
Pull request #429 (2023-10-23 14:52:12)

feat: autofill `model` argument when calling `create_factory` with receiving factory `__model__`
<!--
By submitting this pull request, you agree to:
- follow [Litestar's Code of Conduct](https://github.com/litestar-org/.github/blob/main/CODE_OF_CONDUCT.md)
- follow [Litestar's contribution guidelines](https://github.com/litestar-org/.github/blob/main/CONTRIBUTING.md)
- follow the [PSFs's Code of Conduct](https://www.python.org/psf/conduct/)
-->

### Pull Request Checklist

- [X] New code has 100% test coverage
- [X] (If applicable) The prose documentation has been updated to reflect the changes introduced by this PR
- [X] (If applicable) The reference documentation has been updated to reflect the changes introduced by this PR
- [X] Pre-Commit Checks were ran and passed
- [X] Tests were ran and passed

### Description

This PR aims at allowing to create "sub" factories from concrete factories (i.e: a factory with a `__model__` set) without having to pass the `model` argument to the `create_factory` method.

In this case, when the `model` argument is omitted, the `__model__` of the receiving factory will be used by default, increasing code readability and making the dynamic mode working the same way as the imperative style.

>The PR includes the change, the tests and the documentation update.

```python
from dataclasses import dataclass

from polyfactory.factories import DataclassFactory

@dataclass
class Foo:
    bar: str

# Imperative mode
class FooFactory(DataclassFactory):
    __model__ = Foo

class BarFactory(FooFactory):
    bar = "bar"

# Dynamic mode (after this PR)
foo_factory = DataclassFactory.create_factory(Foo)
bar_factory = foo_factory.create_factory(bar="bar")
```

### Close Issue(s)
Closes #357 

github-actions[bot]: Documentation preview will be available shortly at https://litestar-org.github.io/polyfactory-docs-preview/429
g0di: I made the change as suggested but pyright is not happy. It was the reason why I've made it the previous way. Do you want me to simply turn off the pyright warning for that line or have any suggestion?
guacs: Aah gotcha. I think a pyright ignore is good for that line since I think it's better to have the `AttributeError` as the `__cause__` than do workarounds for satisfying the type checkers.
g0di: > Aah gotcha. I think a pyright ignore is good for that line since I think it's better to have the `AttributeError` as the `__cause__` than do workarounds for satisfying the type checkers.

Fine, I've made the change 🤓 
github-actions[bot]: Documentation preview will be available shortly at https://litestar-org.github.io/polyfactory-docs-preview/429

Patch:
diff --git a/docs/usage/declaring_factories.rst b/docs/usage/declaring_factories.rst
index 6057152..2bfab8c 100644
--- a/docs/usage/declaring_factories.rst
+++ b/docs/usage/declaring_factories.rst
@@ -56,3 +56,11 @@ You can also use this method to create factories imperatively:
 .. literalinclude:: /examples/declaring_factories/test_example_6.py
     :caption: Imperative factory creation
     :language: python
+
+Eventually you can use this method on an existing concrete factory to create a sub factory overriding some parent configuration:
+
+.. literalinclude:: /examples/declaring_factories/test_example_8.py
+    :caption: Imperative sub factory creation
+    :language: python
+
+In this case you don't need to specify the `model` argument to the :meth:`create_factory <polyfactory.factories.base.BaseFactory.create_factory>` method. The one from the parent factory will be used.
diff --git a/polyfactory/factories/base.py b/polyfactory/factories/base.py
index 67ff305..3b8a28c 100644
--- a/polyfactory/factories/base.py
+++ b/polyfactory/factories/base.py
@@ -426,23 +426,30 @@ class BaseFactory(ABC, Generic[T]):
     @classmethod
     def create_factory(
         cls: type[F],
-        model: type[T],
+        model: type[T] | None = None,
         bases: tuple[type[BaseFactory[Any]], ...] | None = None,
         **kwargs: Any,
     ) -> type[F]:
         """Generate a factory for the given type dynamically.
 
-        :param model: A type to model.
+        :param model: A type to model. Defaults to current factory __model__ if any.
+            Otherwise, raise an error
         :param bases: Base classes to use when generating the new class.
         :param kwargs: Any kwargs.
 
         :returns: A 'ModelFactory' subclass.
 
         """
+        if model is None:
+            try:
+                model = cls.__model__
+            except AttributeError as ex:
+                msg = "A 'model' argument is required when creating a new factory from a base one"
+                raise TypeError(msg) from ex
         return cast(
             "Type[F]",
             type(
-                f"{model.__name__}Factory",
+                f"{model.__name__}Factory",  # pyright: ignore[reportOptionalMemberAccess]
                 (*(bases or ()), cls),
                 {"__model__": model, **kwargs},
             ),


Test patch:
diff --git a/docs/examples/declaring_factories/test_example_8.py b/docs/examples/declaring_factories/test_example_8.py
new file mode 100644
index 0000000..8e91f8d
--- /dev/null
+++ b/docs/examples/declaring_factories/test_example_8.py
@@ -0,0 +1,27 @@
+from dataclasses import dataclass
+from enum import Enum
+
+from polyfactory.factories import DataclassFactory
+
+
+class Species(str, Enum):
+    CAT = "Cat"
+    DOG = "Dog"
+    RABBIT = "Rabbit"
+    MOUSE = "Mouse"
+
+
+@dataclass
+class Pet:
+    name: str
+    species: Species
+    sound: str
+
+
+def test_imperative_sub_factory_creation() -> None:
+    pet_factory = DataclassFactory.create_factory(model=Pet)
+    cat_factory = pet_factory.create_factory(species=Species.CAT)
+    cat_instance = cat_factory.build()
+
+    assert isinstance(cat_instance, Pet)
+    assert cat_instance.species == Species.CAT
diff --git a/tests/test_base_factories.py b/tests/test_base_factories.py
index f307bb7..751b6b1 100644
--- a/tests/test_base_factories.py
+++ b/tests/test_base_factories.py
@@ -5,6 +5,7 @@ import pytest
 from pydantic.main import BaseModel
 
 from polyfactory.factories import DataclassFactory
+from polyfactory.factories.base import BaseFactory
 from polyfactory.factories.pydantic_factory import ModelFactory
 
 
@@ -78,3 +79,19 @@ def test_multiple_base_pydantic_factories(override_BaseModel: bool) -> None:
     # see https://github.com/litestar-org/polyfactory/issues/198
     ModelFactory._base_factories.remove(FooModelFactory)
     ModelFactory._base_factories.remove(DummyModelFactory)
+
+
+def test_create_factory_without_model_reuse_current_factory_model() -> None:
+    @dataclass
+    class Foo:
+        pass
+
+    factory = DataclassFactory.create_factory(model=Foo)
+    sub_factory = factory.create_factory()
+
+    assert sub_factory.__model__ == Foo
+
+
+def test_create_factory_from_base_factory_without_providing_a_model_raises_error() -> None:
+    with pytest.raises(TypeError):
+        BaseFactory.create_factory()