Polymorphic subclasses during validation in V2.

[jheddings]

I've been using Pydantic to validate third-party REST API's for a while, and it does a fantastic job keeping the code clean. As I'm moving to V2, there is one area I've struggled to find a suitable replacement: polymorphic subclasses.

I initially implemented the suggestion in #3091 (specifically overriding parse_obj), however that approach does not work for model_validate in V2 (assumedly due to the way it is implemented in pydantic_core). I've also tried following along with suggestions in #5785, however TypeAdapter and discriminated unions are not quite enough for my use case.

Given the following example:

from typing import List, Literal
from pydantic import BaseModel, model_validator

class Animal(BaseModel):
    kind: str

class Pet(Animal):
    name: str

class Cat(Pet):
    kind: Literal['cat']
    meows: int

class Dog(Pet):
    kind: Literal['dog']
    barks: float

class Bird(Pet):
    kind: Literal['bird']
    species: str

class Bee(Animal):
    kind: Literal['bee']
    species: str

class Person(BaseModel):
    name: str
    pets: List[Pet] = []

TIGER = {'name': 'Tiger', 'kind': 'cat', 'meows': 42}
FLUFFY = {'name': 'Fluffy', 'kind': 'dog', 'barks': 3.14}
ACE = {'name': 'Ace', 'kind': 'bird', 'species': 'eagle'}
BUZZ = {'kind': 'bee', 'species': 'western honey bee'}
ALICE = {'name': 'Alice', 'pet': [TIGER, FLUFFY]}  # valid
BOB = {'name': 'Bob', 'pet': [ACE, BUZZ]}  # invalid

As expected, the types for TIGER and FLUFFY are parsed as Pet objects when validated as Person.pet members. Is there a supported method for resolving Person.pet to the correct subclass based on type?

I've tried adding a model_validator to Animal, but that just causes infinite recursion:

    @model_validator(mode="before")
    def make_the_right_animal(cls, data) -> "Animal":
        if data["kind"] == "cat":
            return Cat(**data)
        elif data["kind"] == "dog":
            return Dog(**data)
        elif data["kind"] == "bird":
            return Bird(**data)
        elif data["kind"] == "bee":
            return Bee(**data)
        else:
            raise ValueError(f"Unknown animal kind: {data['kind']}")

It would be possible to list all the valid types rather than the parent class:

class Person(BaseModel):
    name: str
    pets: List[Union[Cat, Dog, Bird]] = []

But this would require a significant refactoring of the codebase and opens up many opportunities for errors (namely when I forget to update lists in each class). I appreciated the clarity of the fully-defined schema when listing out these members, however for this application that would be less than desirable.

It may be "that's just how things work" which is a fine answer... I'm just looking for some direction before I have to find a more "invasive" option.

[sfran96]

My team had the same requirement as #3091, and a solution inspired by the approved answer worked until we upgraded to V2. The inner workings of the new pydantic are somewhat alien to me still, but we managed to get it to work using the documentation for custom data types.

Basically, we can use __get_pydantic_core_schema__ to alter the BaseModel behavior like so.

from typing import Any, Callable, List

from pydantic import BaseModel as BaseModel
from pydantic_core import core_schema


class BaseModel(BaseModel, extra="forbid"):
    __polymorphic__ = False

    @classmethod
    def __get_pydantic_core_schema__(
        cls, __source: type["BaseModel"], __handler: Callable[[Any], core_schema.CoreSchema]
    ) -> core_schema.CoreSchema:
        schema = __handler(__source)
        og_schema_ref = schema["ref"]
        schema["ref"] += ":aux"

        return core_schema.no_info_before_validator_function(
            cls.__convert_to_real_type__, schema=schema, ref=og_schema_ref
        )

    @classmethod
    def __convert_to_real_type__(cls, value: Any):
        if not cls.__polymorphic__:
            return value

        if isinstance(value, dict) is False:
            return value

        value = value.copy()

        subclass = value.pop("kind", None)
        if subclass is None:
            raise ValueError(f"Missing 'kind' in {cls.__name__}")

        try:
            sub = next(dtype for dtype in cls.__subclasses__() if dtype.__name__.lower() == subclass)
        except StopIteration as e:
            raise TypeError(f"Unsupported subclass: {subclass}") from e

        return sub(**value)

    def __init_subclass__(cls, polymorphic: bool = False, **kwargs):
        cls.__polymorphic__ = polymorphic
        super().__init_subclass__(**kwargs)


class Pet(BaseModel, polymorphic=True):
    name: str


class Cat(Pet):
    meows: int


class Dog(Pet):
    barks: float
    friend: Pet


class Bird(Pet):
    species: str


class Bee(Pet):
    species: str


class Person(BaseModel):
    name: str
    pet1: List[Pet] = []
    pet2: List[Pet] = []


pets = []
pets.append({"name": "Tiger", "kind": "cat", "meows": 42})
pets.append({"name": "Fluffy", "kind": "dog", "barks": 3.14, "friend": {"kind": "cat", "meows": 10, "name": "Tiger"}})
pets.append({"name": "Ace", "kind": "bird", "species": "eagle"})
pets.append({"name": "Bee", "kind": "bee", "species": "western honey bee"})

for pet in pets:
    p = Person(**{"name": f"Alice with a {pet['kind']}", "pet1": [pet], "pet2": [pet]})
    print(p)

Pet (and any other class created with polymorphic=True) in this case behaves as an abstract class and can't be instantiated directly, pretty sure that can also be worked around adapting the custom validation function.

The only weird part I don't fully understand why it would be necessary is the following:

        ...
        og_schema_ref = schema["ref"]
        schema["ref"] += ":aux"
        ...

If this is not included, whenever Pet is included more than once in another model, the pre-validation step is only run once, this addresses that.

Similar results can definitely also be achieved using Annotated or BeforeValidator but for our use case we preferred this one.

[hassec]

@sfran96 your link to that issue made me investigate the model_validator and model_serializer possibility, and I think we can get the same behaviour with the simpler class below that avoids __get_pydantic_core_schema__.

class BM(BaseModel, extra="forbid"):
    __is_polymorphic_base = False
    __subclasses_map__ = {}

    @model_serializer(mode="wrap")
    def __serialize_with_class_type__(self, default_serializer) -> Any:
        # default serializer, then append the `__class_name` field and return
        ret = default_serializer(self)
        ret["__class_name"] = f"{self.__class__.__module__}.{self.__class__.__qualname__}"
        return ret

    @model_validator(mode="wrap")
    @classmethod
    def __convert_to_real_type__(cls, value: Any, handler):
        if isinstance(value, dict) is False:
            return handler(value)

        # it is a dict so make sure to remove the __class_name
        # because we don't allow extra keywords but want to ensure
        # e.g Cat.model_validate(cat.model_dump()) works
        class_full_name = value.pop("__class_name", None)

        # if it's not the polymorphic base we construct via default handler
        if not cls.__is_polymorphic_base:
            return handler(value)

        # otherwise we lookup the correct polymorphic type and construct that
        # instead
        if class_full_name is None:
            raise ValueError("Missing __class_name field")

        class_type = cls.__subclasses_map__.get(class_full_name, None)

        if class_type is None:
            # TODO could try dynamic import
            raise TypeError("Trying to instantiate {class_full_name}, which has not yet been defined!")

        return class_type(**value)

    def __init_subclass__(cls, is_polymorphic_base: bool = False, **kwargs):
        cls.__is_polymorphic_base = is_polymorphic_base
        cls.__subclasses_map__[f"{cls.__module__}.{cls.__qualname__}"] = cls
        super().__init_subclass__(**kwargs)

[jheddings]

This is super helpful! I've taken the approach above for deserializing and turned it into a general purpose object. In this approach, objects that extend ABC are ignored and all other subclasses may be considered for validation. This requires setting up the sub-models with enough information to pick the right one. The use of Literal types helps with that.

class AdaptiveObject(BaseModel, ABC):
    """An abstract object that may be one of several concrete types."""

    @model_validator(mode="wrap")
    @classmethod
    def _resolve_adaptive_object(cls, data: dict, handler) -> Any:
        if not isinstance(data, dict):
            return handler(data)

        # if cls is not abstract, there's nothing to do
        if ABC not in cls.__bases__:
            return handler(data)

        # try to validate the data for each possible type
        for subcls in cls._find_all_possible_types():
            try:
                # return the first successful validation
                return subcls.model_validate(data)
            except ValidationError:
                continue

        raise PydanticCustomError(
            "adaptive-object",
            "unable to resolve input",
        )

    @classmethod
    def _find_all_possible_types(cls):
        """Recursively generate all possible types for this object."""

        # any concrete class is a possible type
        if ABC not in cls.__bases__:
            yield cls

        # continue looking for possible types in subclasses
        for subclass in cls.__subclasses__():
            yield from subclass._find_all_possible_types()

class Animal(AdaptiveObject, ABC):
    kind: str

class Pet(Animal, ABC):
    name: str

class Cat(Pet):
    kind: Literal['cat']
    meows: int

...

Now, calling Pet.deserialize(...) will examine all subclasses of Pet to find one that will validate, without having to deal with init_subclass or relying on special fields in the data.

❗ However... I'm stuck on the new serialization behavior. Specifically, the docs indicate that you must use a special annotation to preserve the duck-typing behavior from v1. That's a bummer as it clutters up models and requires altering them to make this work. I have tried modifying the model_serializer but still no luck.

tiger = Cat(**TIGER)
fluffy = Dog(**FLUFFY)
jane = Person(name='Jane', pets=[tiger, fluffy])

#> {'name': 'Jane', 'pets': [{'kind': 'cat', 'name': 'Tiger'}, {'kind': 'dog', 'name': 'Fluffy'}]}

Note the output does not contain all fields for the concrete types. Have you seen this behavior and come up with a workaround?

[hassec]

Sorry, I've just been using SerializeAsAny, so don't know a better solution.

[jheddings]

Looks like we are waiting anxiously for #6423

[hassec]

@sydney-runkle, I noticed that in 2.6.0 the above mentioned option that uses __get_pydantic_core_schema__ does not work anymore.

But it seems that removing the part about the ref

        og_schema_ref = schema["ref"]
        schema["ref"] += ":aux"

and not passing any ref to no_info_before_validator_function, makes the example work again for 2.6.0. Note that this is not going to work for pre 2.6 versions because:

If this is not included, whenever Pet is included more than once in another model, the pre-validation step is only run once, this addresses that.

This isn't urgent as I've moved to the simplified usecase mentioned above.

But I am very curious to better understand what's going on here.
Any idea what changed that could have caused this change in behaviour? I can't really find a lot of information in the docs on what exactly the ref argument is for and how that affects pydantic's behaviour.