Using pydantic to deserialize sublasses of a model

Viewed 1123

I'm using data that follows a class inheritance pattern... I'm having trouble getting pydantic to deserialize it correctly for some use cases.

Given the code below, it appears that the validators are not called when using the parse_* methods. The type for "fluffy" and "tiger" are Animal, however when deserializing the "bob" the Person, his pet is the correct Dog type.

Is there another way to approach this problem? Using pydantic is not a requirement, but the ability to deserialize nested complex types (including List and Dict of objects) is.

# modified from the following examples
# - https://github.com/samuelcolvin/pydantic/issues/2177#issuecomment-739578307
# - https://github.com/samuelcolvin/pydantic/issues/619#issuecomment-713508861

from pydantic import BaseModel

TIGER = """{ "type": "cat", "name": "Tiger the Cat", "color": "tabby" }"""
FLUFFY = """{ "type": "dog", "name": "Fluffy the Dog", "color": "brown", "breed": "rottweiler" }"""

ALICE = """{ "name": "Alice the Person" }"""
BOB = f"""{{ "name": "Bob the Person", "pet": {FLUFFY} }}"""

class Animal(BaseModel):
    type: str
    name: str
    color: str = None

    _subtypes_ = dict()

    def __init_subclass__(cls, type=None):
        cls._subtypes_[type or cls.__name__.lower()] = cls

    @classmethod
    def __get_validators__(cls):
        yield cls._convert_to_real_type_

    @classmethod
    def _convert_to_real_type_(cls, data):
        data_type = data.get("type")

        if data_type is None:
            raise ValueError("Missing 'type' in Animal")

        sub = cls._subtypes_.get(data_type)

        if sub is None:
            raise TypeError(f"Unsupport sub-type: {data_type}")

        return sub(**data)

class Cat(Animal, type="cat"):
    hairless: bool = False

class Dog(Animal, type="dog"):
    breed: str

class Person(BaseModel):
    name: str
    pet: Animal = None

tiger = Animal.parse_raw(TIGER)
print(f"tiger [{tiger.name}] => {type(tiger)} [{tiger.type}]")

fluffy = Animal.parse_raw(FLUFFY)
print(f"fluffy [{fluffy.name}] => {type(fluffy)} [{fluffy.type}]")

bob = Person.parse_raw(BOB)
pet = bob.pet
print(f"bob [{bob.name}] => {type(bob)}")
print(f"pet [{pet.name}] => {type(pet)}")

Output:

tiger [Tiger the Cat] => <class '__main__.Animal'>
fluffy [Fluffy the Dog] => <class '__main__.Animal'>
bob [Bob the Person] => <class '__main__.Person'>
pet [Fluffy the Dog] => <class '__main__.Dog'>
1 Answers

Answer from pydantic project discussion:

You can simply add

class Animal(BaseModel):
    ...
    
    @classmethod
    def parse_obj(cls, obj):
        return cls._convert_to_real_type_(obj)
Related