Repository: litestar-org/polyfactory
Pull request #659 (2025-03-15 22:03:15)

feat: add_providers to easily add new type provider
## Description

Without modifying the factory class. This enables 3rd party custom types to easily add providers.


iloveitaly: Friendly reminder on this pull request. Please provide any updates or additional information if needed.
provinzkraut: @adhtruong?
github-actions[bot]: Documentation preview will be available shortly at https://litestar-org.github.io/polyfactory-docs-preview/659
iloveitaly: Just a friendly reminder on this pull request! Let me know if there's anything else that needs to be addressed.
github-actions[bot]: Documentation preview will be available shortly at https://litestar-org.github.io/polyfactory-docs-preview/659
github-actions[bot]: Documentation preview will be available shortly at https://litestar-org.github.io/polyfactory-docs-preview/659
iloveitaly: Merged your changes in and added a docs example!
github-actions[bot]: Documentation preview will be available shortly at https://litestar-org.github.io/polyfactory-docs-preview/659

Patch:
diff --git a/docs/usage/handling_custom_types.rst b/docs/usage/handling_custom_types.rst
index 42e8712..d8bd705 100644
--- a/docs/usage/handling_custom_types.rst
+++ b/docs/usage/handling_custom_types.rst
@@ -12,6 +12,14 @@ In the above example we override the `get_provider_map` class method, which maps
 returns an appropriate mock value for the type. In the above example, an instance of `CustomSecret` with the hardcoded
 string 'jeronimo'.
 
+Add a Type Provider
+-------------------
+Alternatively, you can add a provider for your custom type to `BaseFactory`.
+
+.. literalinclude:: /examples/handling_custom_types/test_example_3.py
+    :caption: Creating a custom type factory using add_provider
+    :language: python
+
 Creating Custom Base Factories
 ------------------------------
 
diff --git a/polyfactory/factories/base.py b/polyfactory/factories/base.py
index b3608c8..aef5cb8 100644
--- a/polyfactory/factories/base.py
+++ b/polyfactory/factories/base.py
@@ -186,10 +186,14 @@ class BaseFactory(ABC, Generic[T]):
     _factory_type_mapping: ClassVar[dict[Any, type[BaseFactory[Any]]]]
     _base_factories: ClassVar[list[type[BaseFactory[Any]]]]
 
+    _providers: ClassVar[dict[Any, Callable[[], Any]]]
+    """Mapping of type providers that apply to all factories"""
+
     # Non-public attributes
     _extra_providers: dict[Any, Callable[[], Any]] | None = None
+    """Used to copy providers from once base factory to another dynamically generated factory for a class"""
 
-    def __init_subclass__(cls, *args: Any, **kwargs: Any) -> None:  # noqa: C901
+    def __init_subclass__(cls, *args: Any, **kwargs: Any) -> None:  # noqa: C901, PLR0912
         super().__init_subclass__(*args, **kwargs)
 
         if not hasattr(BaseFactory, "_base_factories"):
@@ -198,6 +202,9 @@ class BaseFactory(ABC, Generic[T]):
         if not hasattr(BaseFactory, "_factory_type_mapping"):
             BaseFactory._factory_type_mapping = {}
 
+        if not hasattr(BaseFactory, "_providers"):
+            BaseFactory._providers = {}
+
         if cls.__min_collection_length__ > cls.__max_collection_length__:
             msg = "Minimum collection length shouldn't be greater than maximum collection length"
             raise ConfigurationException(
@@ -417,6 +424,11 @@ class BaseFactory(ABC, Generic[T]):
 
     # Public Methods
 
+    @classmethod
+    def add_provider(cls, provider_type: Any, provider_function: Callable[[], Any]) -> None:
+        """Add a provider for a custom type to be available to all factories"""
+        cls._providers[provider_type] = provider_function
+
     @classmethod
     def is_factory_type(cls, annotation: Any) -> bool:
         """Determine whether a given field is annotated with a type that is supported by a base factory.
@@ -502,7 +514,7 @@ class BaseFactory(ABC, Generic[T]):
 
     @classmethod
     def get_provider_map(cls) -> dict[Any, Callable[[], Any]]:
-        """Map types to callables.
+        """Map types to callables which accept no arguments and return a random value of the given type.
 
         :notes:
             - This method is distinct to allow overriding.
@@ -552,6 +564,7 @@ class BaseFactory(ABC, Generic[T]):
             Callable: _create_generic_fn,
             abc.Callable: _create_generic_fn,
             Counter: lambda: Counter(cls.__faker__.pystr()),
+            **(cls._providers or {}),
             **(cls._extra_providers or {}),
         }
 
@@ -834,7 +847,11 @@ class BaseFactory(ABC, Generic[T]):
             with suppress(Exception):
                 return unwrapped_annotation()
 
-        msg = f"Unsupported type: {unwrapped_annotation!r}\n\nEither extend the providers map or add a factory function for this type."
+        msg = (
+            f"Unsupported type: {unwrapped_annotation!r} on field '{field_meta.name}' from class {cls.__name__}."
+            "\n\nEither use 'add_provider', extend the providers map, or add a factory function for the field on the model."
+        )
+
         raise ParameterException(
             msg,
         )
@@ -1126,7 +1143,7 @@ class BaseFactory(ABC, Generic[T]):
 
     @classmethod
     def coverage(cls, **kwargs: Any) -> abc.Iterator[T]:
-        """Build a batch of the factory's Meta.model will full coverage of the sub-types of the model.
+        """Build a batch of the factory's Meta.model with full coverage of the sub-types of the model.
 
         :param kwargs: Any kwargs. If field_meta names are set in kwargs, their values will be used.
 


Test patch:
diff --git a/docs/examples/handling_custom_types/test_example_3.py b/docs/examples/handling_custom_types/test_example_3.py
new file mode 100644
index 0000000..72e0c53
--- /dev/null
+++ b/docs/examples/handling_custom_types/test_example_3.py
@@ -0,0 +1,43 @@
+from dataclasses import dataclass
+from typing import TypeVar
+from uuid import UUID
+
+from polyfactory.factories import DataclassFactory
+from polyfactory.factories.base import BaseFactory
+
+
+# we created a special class we will use in our code
+class CustomSecret:
+    def __init__(self, value: str) -> None:
+        self.value = value
+
+    def __repr__(self) -> str:
+        return "*" * len(self.value)
+
+    def __str__(self) -> str:
+        return "*" * len(self.value)
+
+
+T = TypeVar("T")
+
+
+# we add a function to generate the custom type
+BaseFactory.add_provider(CustomSecret, lambda: CustomSecret("jeronimo"))
+
+
+@dataclass
+class Person:
+    id: UUID
+    secret: CustomSecret
+
+
+# we don't have to use a custom base factory!
+class PersonFactory(DataclassFactory[Person]): ...
+
+
+def test_custom_dataclass_base_factory() -> None:
+    person_instance = PersonFactory.build()
+    assert isinstance(person_instance.secret, CustomSecret)
+    assert repr(person_instance.secret) == "*" * len("jeronimo")
+    assert str(person_instance.secret) == "*" * len("jeronimo")
+    assert person_instance.secret.value == "jeronimo"
diff --git a/tests/test_provider_map.py b/tests/test_provider_map.py
index 75c54eb..9a6096c 100644
--- a/tests/test_provider_map.py
+++ b/tests/test_provider_map.py
@@ -1,6 +1,9 @@
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, Generic, TypeVar
 
+import pytest
+
+from polyfactory.exceptions import ParameterException
 from polyfactory.factories.base import BaseFactory
 from polyfactory.factories.dataclass_factory import DataclassFactory
 
@@ -55,3 +58,23 @@ def test_provider_map_with_typevar() -> None:
 
     coverage_result = list(FooFactory.coverage())
     assert all(result.foo == "any" for result in coverage_result)
+
+
+def test_add_custom_provider() -> None:
+    class CustomType:
+        def __init__(self, _: Any) -> None:
+            pass
+
+    @dataclass
+    class Foo:
+        foo: CustomType
+
+    FooFactory = DataclassFactory.create_factory(Foo)
+
+    with pytest.raises(ParameterException):
+        FooFactory.build()
+
+    BaseFactory.add_provider(CustomType, lambda: CustomType("custom"))
+
+    # after adding the provider, nothing should raise!
+    assert FooFactory.build()