How to make a pydantic Field accept subclasses using Type?

Question

I'm trying to have a field in one Pydantic model accept any of a set of BaseModel-derived classes or subclasses that I define separately. Reading the docs here, I naively did the below, which failed; I then realised that I'd misread the docs and that in this scenario "a field may only accept classes (not instances)", and also that Foo and Bar in that example don't derive from BaseModel themselves (is that important?).

I'm guessing that I'm just misconceiving this from the start, so my question: is there a correct way to do what I'm trying to do without using a Union on the subclasses, or some other way entirely that's better?

Bonus question: what's a common usecase for only being able to accept classes and not instances?

MRE:

from pydantic import BaseModel

from typing import Type, Union

class Foo(BaseModel):
    pass

class Bar(Foo):
    pass

class Baz(Foo):
    pass

class Container(BaseModel):
    some_foo: Type[Foo] # this fails
    # this will run successfully --> some_foo: Union[Bar, Baz]


b = Baz()
c = Container(some_foo = b)
# Traceback (most recent call last):
#   File "mre.py", line 20, in <module>
#     c = Container(some_foo = b)
#   File "pydantic/main.py", line 400, in pydantic.main.BaseModel.__init__
# pydantic.error_wrappers.ValidationError: 1 validation error for Container
# some_foo
#   subclass of Foo expected (type=type_error.subclass; expected_class=Foo)

Is something works wrong if you define attribute as `some_foo: Foo`? Question is a bit confusing, you set restriction as Type, but then create class with instance. — NobbyNobbs, May 03 '21 at 11:41
@NobbyNobbs You're right, I should have been clearer. The problem with `some_foo: Foo` is that it doesn' validate properly (which @p3j4p5's answer picked up on brilliantly). If put `some_foo: Foo`, you can put pretty much any class instance in and it will be accepted (including, say, `class NotFoo(BaseModel): pass`. — Winawer, May 03 '21 at 21:08

Paul P · Accepted Answer · 2021-05-03T13:07:34.223

This is trickier than it seems.

Here is a solution that works using pydantic's validator but maybe there is a more "pydantic" approach to it.

from pydantic import BaseModel, validator
from typing import Any

class Foo(BaseModel):
    pass

class Bar(Foo):
    pass

class Baz(Foo):
    pass

class NotFoo(BaseModel):
    pass

class Container(BaseModel):
    some_foo: Any

    @validator("some_foo")
    def validate_some_foo(cls, val):
        if issubclass(type(val), Foo):
            return val

        raise TypeError("Wrong type for 'some_foo', must be subclass of Foo")

b = Bar()
c = Container(some_foo=b)
# Container(some_foo=Bar())

n = NotFoo()
c = container(some_foo=n)

# Traceback (most recent call last):
#   File "/path/to/file.py", line 64, in <module>
#     c = Container(some_foo=n)
#   File "pydantic/main.py", line 400, in pydantic.main.BaseModel.__init__
# pydantic.error_wrappers.ValidationError: 1 validation error for Container
# some_foo
#   Wrong type for 'some_foo', must be subclass of Foo (type=type_error)

Please note that the custom validator must raise a ValueError or TypeError (or sub-classes thereof) for pydantic to be able to properly re-raise a ValudationError.

One reason why you might want to have a specific class (as opposed to an instance of that class) as the field type is when you want to use that field to instantiate something later on using that field.

Here's an example:

from pydantic import BaseModel
from typing import Optional, Type

class Foo(BaseModel):
    # x is NOT optional
    x: int

class Bar(Foo):
    y: Optional[str]

class Baz(Foo):
    z: Optional[bool]

class NotFoo(BaseModel):
    # a is NOT optional
    a: str

class ContainerForClass(BaseModel):
    some_foo_class: Type[Foo]

c = ContainerForClass(some_foo_class=Bar)


# At this point you know that you will use this class for something else
# and that x must be always provided and it must be an int:
d = c.some_foo_class(x=5, y="some string")
# Baz(x=5, z=None)


c = ContainerForClass(some_foo_class=Baz)

# Same here with x:
e = c.some_foo_class(x=6, z=True)
# Baz(x=6, z=True)


# Would't work with this:
c = ContainerForClass(some_foo_class=NotFoo)

# Traceback (most recent call last):
#   File "/path/to/file.py", line 98, in <module>
#     c = ContainerForClass(some_foo_class=NotFoo)
#   File "pydantic/main.py", line 400, in pydantic.main.BaseModel.__init__
# pydantic.error_wrappers.ValidationError: 1 validation error for ContainerForClass
# some_foo_class
#   subclass of Foo expected (type=type_error.subclass; expected_class=Foo

I changed the `some_foo` type to `Foo` rather than `Any`. Otherwise `mypy` will allow something like `c = Container(some_foo="string")`. Any reason that `Any` would be preferable? (I'm sure there are still some quirks of Python's type annotations I'm unaware of). — Matt Shirley, Jul 04 '23 at 20:12

How to make a pydantic Field accept subclasses using Type?

1 Answers1