Pydantic Serialization of Complex Keys in Dicts/Maps

example.py

"""
https://github.com/pydantic/pydantic/issues/10724

Example of using complex keys in Dict using Pydantic for Serialization

There's a two different approaches
1. Fundamentally hook into the core serialization and map the structure into a list[{Key, Value}] (or similar) to get the Obj -> Json -> Obj working
2. Create an intermediate Object in Pydantic and add `to_X` method and `from_X` classmethod to get the Obj -> Obj2 -> Json -> Obj2 -> Obj

Method 2 is shown below
"""
from typing import Generic, TypeVar, Self

from pydantic import BaseModel, ConfigDict

K = TypeVar("K")
V = TypeVar("V")


class KeyValue(BaseModel, Generic[K, V]):
    key: K
    value: V


class Alpha(BaseModel):
    model_config = ConfigDict(frozen=True)
    name: str
    counter: int


class Beta(BaseModel):
    a: str
    b: dict[Alpha, str]


class BetaIntermediate(BaseModel):
    a: str
    b: list[KeyValue[Alpha, str]]

    def to_beta(self) -> Beta:
        return Beta(a=self.a, b={x.key: x.value for x in self.b})

    @classmethod
    def from_beta(cls, beta: Beta) -> Self:
        return BetaIntermediate(
            a=beta.a,
            b=[KeyValue[Alpha, str](key=k, value=v) for k, v in beta.b.items()],
        )


def demo() -> None:

    beta = Beta(
        a="shrimp",
        b={Alpha(name="a0", counter=7): "0", Alpha(name="a1", counter=1): "1"},
    )

    # intermediate structure
    i = BetaIntermediate.from_beta(beta)
    jx = i.model_dump_json()
    beta2 = i.model_validate_json(jx).to_beta()

    print((beta, beta2))
    print(beta == beta2)


if __name__ == "__main__":
    demo()