[Python] Implementing Discriminated Unions in Pydantic without using Nested Models?

I'm trying to implement discriminated unions in Pydantic to select the correct class based on user input using the discriminator parameter. While the documentation suggests creating a nested model class to handle this easily, I'd like to use this functionality without introducing an additional nested model and have a similar behaviour as a normal Pydantic BaseModel class.

I've tried to use RootModel as a workaround, but the resulting object is encapsulated within the .root property, which isn't ideal for my use case. I am able to do .model_dump() but unable to access the attributes on it directly.

Is there a better way to implement this without creating a nested model or using RootModel?

from typing import Literal, Union, Annotated

from pydantic import BaseModel, Field, RootModel


class Cat(BaseModel):
pet_type: Literal["cat"]
meows: int


class Dog(BaseModel):
pet_type: Literal["dog"]
barks: float


class Lizard(BaseModel):
pet_type: Literal["reptile", "lizard"]
scales: bool


Animal = Annotated[
Union[Cat, Dog, Lizard],
Field(discriminator="pet_type"),
]
AnimalModel = RootModel[Animal]

animal = AnimalModel.model_validate({"pet_type": "cat", "meows": 3})

try:
# want to access the attributes directly
print(animal.pet_type)
except AttributeError as e:
print(e)
#> "RootModel[Annotated[Union[Cat, Dog, Lizard], FieldInfo(annotation=NoneType, required=True, discriminator='pet_type')]]" object has no attribute 'pet_type'

# have to access the attributes by first accessing the .root attribute
print(animal.root.pet_type)


(am using Pydantic v2.7)

Continue reading...