blog.notmyidea.org/content/code/2024-07-12-pydantic.md

title	tags
Discriminate pydantic objects by field	pydantic, python, match

I really like Pydantic because it makes it easy to define the the structure of the objects I want to use, using typing.

I wanted to parse a json object and build a different object depending on the value of some specific key in the JSON object.

Here, I have three types of messages: OperationMessage, PeerMessage and ServerRequest, as follows:

from typing import Literal, Optional
from pydantic import BaseModel


class OperationMessage(BaseModel):
    kind: Literal["operation"] = "operation"
    verb: Literal["upsert", "update", "delete"]


class PeerMessage(BaseModel):
    kind: Literal["peermessage"] = "peermessage"
    sender: str
    recipient: str
    message: dict


class ServerMessage(BaseModel):
    kind: Literal["server"] = "server"
    action: Literal["list-peers"]

Each of these classes share the same kind key, but each of them has a different value for it. it's our discriminator.

Let's build a generic Request class that will be able to build for me the proper objects:

from typing import Union
from pydantic import Field, RootModel, ValidationError


class Request(RootModel):
    root: Union[ServerMessage, PeerMessage, OperationMessage] = Field(
        discriminator="kind"
    )

Which can be used this way:

try:
    incoming = Request.model_validate_json(raw_message)
except ValidationError as e:
    # Oh noes.

Because we have classes, we can leverage the match statement:

match incoming.root:
    case OperationMessage():
        # Here to broadcast the message
        websockets.broadcast(peers, raw_message)

    case PeerMessage(recipient=_id):
        # Or to send peer messages to the proper peer
        peer = connections.get(_id)
        if peer:
            await peer.send(raw_message)
    # ... Etc.