I have to model a pydantic class from a JSON object which contains some invalid syntax keys.
As an example:
example = {
"$type": "Menu",
"name": "lunch",
"children": [
{"$type": "Pasta", "title": "carbonara"},
{"$type": "Meat", "is_vegetable": false},
]
}
My pydantic classes at the moment looks like:
class Pasta(BaseModel):
title: str
class Meat(BaseModel):
is_vegetable: bool
class Menu(BaseModel):
name: str
children: list[Pasta | Meat]
Now, this work except for $type
field. If the field was called "dollar_type", I would simply create the following TranslationModel
base class and let Pasta
, Meat
and Menu
inherit from TranslationModel
:
class TranslationModel(BaseModel):
@computed_field
def dollar_type(self) -> str:
return self.__class__.__name__
so that by executing Menu(**example).model_dump()
I get
{
'dollar_type': 'Menu',
'name': 'lunch',
'children': [
{'dollar_type': 'Pasta', 'title': 'carbonara'},
{'dollar_type': 'Meat', 'is_vegetable': False}
]
}
But sadly I have to strictly follow the original json structure, so I have to use $type
.
I have tried using alias
and model_validator
by following the documentation but without success.
How could I solve this? Thanks in advance
This very much looks like you would rather apply a discriminated union pattern. See the following example:
from pydantic import BaseModel, Field
from typing import Literal, Annotated
example = {
"$type": "Menu",
"name": "lunch",
"children": [
{"$type": "Pasta", "title": "carbonara"},
{"$type": "Meat", "is_vegetable": False},
]
}
class Pasta(BaseModel):
type: Literal["Pasta"] = Field("Pasta", alias="$type")
title: str
class Meat(BaseModel):
type: Literal["Meat"] = Field("Meat", alias="$type")
is_vegetable: bool
AnyDish = Annotated[Pasta | Meat, Field(discriminator="type")]
class Menu(BaseModel):
name: str
children: list[AnyDish]
menu = Menu.model_validate(example)
print(menu)
menu.model_dump(by_alias=True)
Which prints:
name='lunch' children=[Pasta(type='Pasta', title='carbonara'), Meat(type='Meat', is_vegetable=False)]
{'name': 'lunch', 'children': [{'$type': 'Pasta', 'title': 'carbonara'}, {'$type': 'Meat', 'is_vegetable': False}]}
This pattern has multiple advantages:
AnyDish
type.To reduce verbosity one can also introduce a short-cut like:
class TypeLiteral:
def __class_getitem__(cls, tag: str):
return Annotated[Literal[tag], Field(default=tag, alias="$type")]
class Pasta(BaseModel):
type: TypeLiteral["Pasta"]
title: str
class Meat(BaseModel):
type: TypeLiteral["Meat"]
is_vegetable: bool
You can find more about discriminated unions in the pydantic docs: https://docs.pydantic.dev/latest/concepts/unions/#discriminated-unions
I hope this is useful!