Search code examples
pythonpydanticpydantic-v2

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


Solution

  • 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!