litestar-org/polyfactory

Bug: TypeAliases/NewTypes ignored in Provider Map

issacruthe opened this issue · 4 comments

Description

I discovered this bug when using pydantic, which makes heavy use of annotated types and type aliases when defining custom validators for types. From their documentation:

from typing import Annotated

from pydantic import BaseModel, AfterValidator, ValidationError


MyInt = Annotated[int, AfterValidator(lambda v: v + 1)]

class Model(BaseModel):
    a: MyInt

print(Model(a=1).a)
# > 2

From what I understand, polyfactory lets the user define a custom provider map to override the default handling of user-defined types. Drilling down into polyfactory's code, I've found that the provider_map does not support "unwrap-able" types because all annotations, newtypes, optionals, and unions are removed or "unwrapped" before any attempt is made to find a type in the provider map. This means even if I add MyInt to the provider_map with a custom factory method, polyfactory will always treat MyInt as just an int because Annotated will be stripped by the unwrap_annotation method.

URL to code causing the issue

unwrapped_annotation = unwrap_annotation(field_meta.annotation, random=cls.__random__)

field_type = normalize_annotation(annotation, random=random)

MCVE

from dataclasses import dataclass
from typing import NewType
from polyfactory.factories import DataclassFactory

MyType = NewType("MyType", str)


@dataclass
class MyModel:
    my_type: MyType


class MyModelFactory(DataclassFactory):
    __model__ = MyModel

    @classmethod
    def get_provider_map(cls):
        base_provider_map = super().get_provider_map()
        return {
            MyType: lambda: "constant",  # Any MyType value *should* always be populated with "constant"
            **base_provider_map
        }


if __name__ == "__main__":
    my_model = MyModelFactory.build()
    print(my_model.my_type)

Results in:

# stdout
ddtJwRSSJdmQZPLLcuuq  # expected value is "constant"

Release Version

2.8.0

Platform

  • Linux
  • Mac
  • Windows
  • Other (Please specify in the description above)

Funding

  • If you would like to see an issue prioritized, make a pledge towards it!
  • We receive the pledge once the issue is completed & verified
Fund with Polar

I've taken a brief stab at fixing this with the following code, which basically hijacks the unwrap_annotation function. However, I've found that even at this point the annotation has already gone through a couple passes of normalization of the NewType and Optional types, particularly when initializing the FieldMeta object, so this code works marginally better, but only for a subset of (more complex) custom types.

from dataclasses import dataclass, asdict
from typing import Any, Generator, Annotated, Optional, NewType
from polyfactory.factories import DataclassFactory
from polyfactory.field_meta import FieldMeta


MyType1 = NewType("MyType1", str)
MyType2 = Optional[Annotated[str, "metadata"]]


@dataclass
class MyModel:
    my_type1: MyType1
    my_type2: MyType2


class MyModelFactory(DataclassFactory):
    __model__ = MyModel

    @classmethod
    def get_provider_map(cls):
        base_provider_map = super().get_provider_map()
        return {
            MyType1: lambda: "constant1",  # Any MyType1 value *should* always be populated with "constant1"
            MyType2: lambda: "constant2",  # Any MyType2 value *should* always be populated with "constant2"
            **base_provider_map
        }

    @classmethod
    def get_field_value(cls, field_meta: FieldMeta, field_build_parameters: Any = None) -> Any:
        for annotation in cls.unwrap_annotation(field_meta.annotation):
            if provider := cls.get_provider_map().get(annotation):
                return provider()
        return super().get_field_value(field_meta, field_build_parameters)

    @classmethod
    def unwrap_annotation(cls, annotation: Any) -> Generator:
        # There is a bug in the polyfactory library where type detection for TypeAliases and NewTypes drill down all
        # the way to the base type before checking if the user provided a factory method for that type in the provider
        # map. This custom loop builds from their existing logic for unwrapping annotations in
        # polyfactory.utils.helpers.unwrap_annotation and exposes every intermediary type so that we can check
        # that type against the user-specified provider map.
        from polyfactory.utils.helpers import (
            is_optional, is_union, is_annotated, is_new_type
        )
        from typing_extensions import get_args
        while is_optional(annotation) or is_union(annotation) or is_new_type(annotation) or is_annotated(annotation):
            yield annotation
            if is_new_type(annotation):
                annotation = annotation.__supertype__
            elif is_optional(annotation) or is_annotated(annotation):
                annotation = next((arg for arg in get_args(annotation) if arg not in (type(None), None)))
            elif is_union(annotation):
                args = list(get_args(annotation))
                annotation = cls.__random__.choice(args)
            else:
                break


if __name__ == "__main__":
    my_model = MyModelFactory.build()
    print(asdict(my_model))
# stdout
{'my_type1': 'zvivkqTnOUXPsWRnYEzZ', 'my_type2': 'constant2'}

So this partial fix addresses my_type_2, but not my_type_1

guacs commented

Thanks for reporting this :)

I agree that overriding the provider map for a NewType should be properly considered since that is, for all practical purposes, a new type. However, I'm not so sure about the type aliases. In my opinion, type aliases are not user defined types, and the provider map is for handling cases of user defined types that polyfactoey can't create instances of.

@Goldziher what do you think about this?

Thanks for reporting this :)

I agree that overriding the provider map for a NewType should be properly considered since that is, for all practical purposes, a new type. However, I'm not so sure about the type aliases. In my opinion, type aliases are not user defined types, and the provider map is for handling cases of user defined types that polyfactoey can't create instances of.

@Goldziher what do you think about this?

i agree. A new type is a new type, and we have some logic to handle it. We obviously should allow it being overridden as in the bug issue.

As for type aliases - these are not actually distinct types, and we cannot handle this type of logic like this. I would say this goes onto the user to handle. Also note-- if you actually want to have a type alias idiomatically, you should type it using TypeAlias exported from typing.

Thanks for reporting this :)

I agree that overriding the provider map for a NewType should be properly considered since that is, for all practical purposes, a new type. However, I'm not so sure about the type aliases. In my opinion, type aliases are not user defined types, and the provider map is for handling cases of user defined types that polyfactoey can't create instances of.

@Goldziher what do you think about this?

Thank you both for your timely replies - I see your points and think supporting NewTypes in the provider map would unlock the functionality I’m after and is a reasonable compromise. I also see how supporting type aliases in the provider map could have unintended consequences and is best steered clear of.