Handling Custom Types

Sometimes you need to handle custom types, either from 3rd party libraries or from your own codebase. To achieve this, you can extend what’s referred as the providers_map:

Extending the provider map

from dataclasses import dataclass
from typing import Any, Dict, Type
from uuid import UUID

from polyfactory.factories import DataclassFactory


# 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)


@dataclass
class Person:
    id: UUID
    secret: CustomSecret


# by default the factory class cannot handle unknown types,
# so we need to override the provider map to add it:
class PersonFactory(DataclassFactory[Person]):
    @classmethod
    def get_provider_map(cls) -> Dict[Type, Any]:
        providers_map = super().get_provider_map()

        return {
            CustomSecret: lambda: CustomSecret("jeronimo"),
            **providers_map,
        }


def test_custom_secret_creation() -> 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"

In the above example we override the get_provider_map class method, which maps types to callables. Each callable in the map returns an appropriate mock value for the type. In the above example, an instance of CustomSecret with the hardcoded string ‘jeronimo’.

Creating Custom Base Factories

The above works great when you need to do this in a localised fashion, but if you need to use some custom types in many places it will lead to unnecessary duplication. The solution for this is to create a custom base factory, in this case for handling dataclasses:

Creating a custom dataclass factory with extended provider map

from dataclasses import dataclass
from typing import Any, Dict, Generic, Type, TypeVar
from uuid import UUID

from polyfactory.factories import DataclassFactory


# 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 create a custom base factory to handle dataclasses, with an extended provider map
class CustomDataclassFactory(Generic[T], DataclassFactory[T]):
    __is_base_factory__ = True

    @classmethod
    def get_provider_map(cls) -> Dict[Type, Any]:
        providers_map = super().get_provider_map()

        return {
            CustomSecret: lambda: CustomSecret("jeronimo"),
            **providers_map,
        }


@dataclass
class Person:
    id: UUID
    secret: CustomSecret


# we use our CustomDataclassFactory as a base for the PersonFactory
class PersonFactory(CustomDataclassFactory[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"

Note

If extra configs values are defined for custom base classes, then __config_keys__ should be extended so that these values are correctly passed onto to concrete factories.