from pydantic import BaseModel, ConfigDict
from polyfactory.factories.pydantic_factory import ModelFactory

class CustomSecretPydantic(BaseModel):
    value: str


class CustomSecretPlain:
    def __init__(self, value: str) -> None:
        self.value = value

    def __repr__(self) -> str:
        return self.value

    def __str__(self) -> str:
        return self.value


class MyModel(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    id: int
    secret: CustomSecretPydantic


class MyFactory(ModelFactory[MyModel]):

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

        return {
            **providers_map,
            CustomSecretPydantic: lambda: CustomSecretPydantic.model_validate({"value": "asdf"}),
            CustomSecretPlain: lambda: CustomSecretPlain(value="asdf"),
        }


if __name__ == "__main__":

    print(MyFactory.build())