Using Discriminators in Pydantic

pydantic_discriminators.py

import doctest
import typing as t
from decimal import Decimal

from pydantic import BaseModel, Field


class SadClown(BaseModel):
    mood: t.Literal["😟"]
    salary: Decimal = Field(..., lt=1_000_000_000.00)


class HappyClown(BaseModel):
    mood: t.Literal["😁"]
    salary: Decimal = Field(..., gte=1_000_000_000.00)


Clown = t.Annotated[t.Union[SadClown, HappyClown], Field(discriminator="mood")]


class Performance(BaseModel):
    """
    Example of JSON Schema / Pydantic discriminators
    ================================================

    .. note::
        For this to work out of the box, you'll need to have an environment
        and Pydantic installed.  A quick way to run this w/o making a huge
        mess:

        .. code-block:: shell

            python3 -m venv .venv
            source .venv/bin/activate
            python -m pydantic_example.py

            # or you can copy this entire file's content to your clipboard
            # python -c "$(pbpaste)"

    Take this example payload, say from a web service (raw JSON)...
    >>> performance_response = '{"actor": {"mood":"😟", "salary": 50.18, "shoe_size": "large"}}'

    Now, we deserialize it:
    >>> Performance.parse_raw(performance_response.encode('utf-8'))
    Performance(actor=SadClown(mood='😟', salary=Decimal('50.18')))

    Here's what the JSON Schema looks like:
    >>> print(Performance.schema_json(indent=2))  # doctest: +ELLIPSIS,+NORMALIZE_WHITESPACE
    {
        "title": "Performance",
            ...

        "properties": {
            "actor": {
                "title": "Actor",
                "discriminator": {
                    "propertyName": "mood",
                    "mapping": {
                        "😟": "#/definitions/SadClown",
                        "😁": "#/definitions/HappyClown"
                    }
                },
            "oneOf": [
                {
                  "$ref": "#/definitions/SadClown"
                },
                {
                  "$ref": "#/definitions/HappyClown"
                }
            ]
        }
      },
      "required": [
        "actor"
      ],
      "definitions": {
        "SadClown": {
          "title": "SadClown",
          "type": "object",
          "properties": {
            "mood": {
              "title": "Mood",
              "enum": [
                "😟"
              ],
              "type": "string"
            },
            "salary": {
              "title": "Salary",
              "exclusiveMaximum": 1000000000.0,
              "type": "number"
            }
          },
          "required": [
            "mood",
            "salary"
          ]
        },
        "HappyClown": {
          "title": "HappyClown",
          "type": "object",
          "properties": {
            "mood": {
              "title": "Mood",
              "enum": [
                "😁"
              ],
              "type": "string"
            },
            "salary": {
              "title": "Salary",
              "gte": 1000000000.0,
              "type": "number"
            }
          },
          "required": [
            "mood",
            "salary"
          ]
        }
      }
    }
    """

    actor: Clown


doctest.testmod(verbose=True)