Pydantic: Is it appropriate to implement a poly model using validate func?

[binderclip]

sample code:

from enum import IntEnum

from pydantic import BaseModel


class FooTypeEnum(IntEnum):
    A = 1
    B = 2


class Foo(BaseModel):
    type: int

    @classmethod
    def validate(cls, value):
        type_to_model = {
            FooTypeEnum.A: FooA,
            FooTypeEnum.B: FooB,
            None: cls
        }
        _type = value.get('type', None)
        try:
            return type_to_model[_type](**value)
        except KeyError:
            raise TypeError(f'no model for type: {_type}')


class FooA(Foo):
    a: str


class FooB(Foo):
    b: str


class Bar(BaseModel):
    foo: Foo = ...


m = Bar(foo={'type': FooTypeEnum.A, 'a': 'aaa'})
print(m)
# Bar foo=<FooA type=1 a='aaa'>
print(m.dict())
# {'foo': {'type': 1, 'a': 'aaa'}}
m = Bar(foo={'type': FooTypeEnum.B, 'b': 'bbb'})
print(m)
# Bar foo=<FooB type=2 b='bbb'>
print(m.dict())
# {'foo': {'type': 2, 'b': 'bbb'}}

# m = Bar(foo={'type': FooTypeEnum.B, 'c': 'bbb'})
# error ...
# m = Bar(foo={'type': 3, 'c': 'bbb'})
# error ...
# m = Bar(foo={'c': 'bbb'})
# error ...

[samuelcolvin]

I'm really not sure what your question is here. What does "poly model" mean?

[binderclip]

In schematics, there is a PolyModelType.

"""A field that accepts an instance of any of the specified models."""

https://github.com/schematics/schematics/blob/master/schematics/types/compound.py#L333

And there is a example using PolyModelType to implement a Msg model (which could be TextMsg or ImageMsg): https://github.com/binderclip/code-snippets-python/blob/master/models/schematics_snippets/schematics_poly_model_type.py#L55-L81

I want to use pydantic to do the same thing.

[samuelcolvin]

Looks like you just need to use Union:

from typing import Union
from enum import IntEnum

from pydantic import BaseModel


class FooTypeEnum(IntEnum):
    A = 1
    B = 2


class Foo(BaseModel):
    type: FooTypeEnum


class FooA(Foo):
    a: str


class FooB(Foo):
    b: str


class Bar(BaseModel):
    foo: Union[FooA, FooB]


m = Bar(foo={'type': FooTypeEnum.A, 'a': 'aaa'})
print(m)
# Bar foo=<FooA type=1 a='aaa'>
print(m.dict())
# {'foo': {'type': 1, 'a': 'aaa'}}
m = Bar(foo={'type': FooTypeEnum.B, 'b': 'bbb'})
print(m)
# Bar foo=<FooB type=2 b='bbb'>
print(m.dict())
# {'foo': {'type': 2, 'b': 'bbb'}}

# m = Bar(foo={'type': FooTypeEnum.B, 'c': 'bbb'})
# error ...
# m = Bar(foo={'type': 3, 'c': 'bbb'})
# error ...
# m = Bar(foo={'c': 'bbb'})
# error ...

Pydantic tries each of the different types in Union to see if they're valid and returns the first which is.

[binderclip]

Union can do the job even easier, thank you!

[samuelcolvin]

no problem, glad it helped.