python pydantic pattern matching

pydantic_pattern_match.py

from dataclasses import dataclass
from enum import Enum, auto
from typing import Annotated, Literal
from pydantic import BaseModel, Field


class VarType(str, Enum):
    a = "a"
    b = "b"


class VarA(BaseModel):
    var_type: Literal[VarType.a] = VarType.a
    field_a: int
    some_other_field: str


class VarB(BaseModel):
    var_type: Literal[VarType.b] = VarType.b
    field_b: int


Var = VarA | VarB


class Box(BaseModel):
    var: Var = Annotated[..., Field(discriminator="var_type")]


def show_var(v: Var):
    match v:
        case VarA(some_other_field=f):
            print(f"Got `VarA` `some_other_field` = `{f}`")
        case VarB():
            print(f"Got `VarB`")


show_var(VarA(field_a=3, some_other_field="asd"))
show_var(VarB(field_b=3))

box = Box.model_validate(
    {
        "var": dict(
            var_type="b",
            field_b=2,
        )
    }
)
show_var(box.var)