CODEWITHSUNDEEP
pythonpydanticpydantic-v2
ndricca
ndricca

Reputation: 532

create pydantic computed field with invalid syntax name

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

Upvotes: 2

Views: 315

Answers (2)

Axel Donath
Axel Donath

Reputation: 1638

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:

  • It decouples the class name from the tag. This is usually preferred, because class names might change. But you might still want to read old files.
  • The pattern is extensible and explicit. You can easily add new types of dishes later and just include them in the 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!

Upvotes: 0

ndricca
ndricca

Reputation: 532

Obviously I wrote the question after more than 1 hour digging around this, and just after posting it I found a solution by myself:

class TranslationModel(BaseModel):

    @computed_field(alias="$type")
    def dollar_type(self) -> str:
        return self.__class__.__name__

class Pasta(TranslationModel):
    title: str

class Meat(TranslationModel):
    is_vegetable: bool

class Menu(TranslationModel):
    name: str
    children: list[Pasta | Meat]

I already tried setting alias="$type" into computed_field decorator, what I was missing was calling model_dump with by_alias=True:

Menu(**example).model_dump(by_alias=True)

{'name': 'lunch', 'children': [{'title': 'carbonara', '$type': 'Pasta'}, {'is_vegetable': False, '$type': 'Meat'}], '$type': 'Menu'}

Upvotes: 1

Related Questions