Custom Dictionary Types in Pydantic

In one of my projects at work, I wanted to define a custom dictionary type using Pydantic. The aim was to define the keys using an Enum class and values using a model type.

It wasn’t obvious at first how to approach this task. But, after some troubleshooting, I found a reasonable long-term solution.

Contents

• The Problem
• Brute-force Approach
• Final Solution
• Conclusion

The Problem

I had already defined the Enum and model types, which I used throughout the codebase. For example, the types looked like these:

from enum import Enum
from pydantic import BaseModel

class AnimalSpecies(str, Enum):
	LION = "lion"
	DOLPHIN = "dolphin"
	ZEBRA = "zebra"

class AnimalData(BaseModel):
	habitat: str
	diet: str

Here, we have an Enum type containing different animal species and a model containing data about a given species.

What if we want to create a dictionary with keys defined by AnimalSpecies and values defined by AnimalData?

Brute-force Approach

A quick way to do this would be to write

class AnimalDict(BaseModel):
	lion: AnimalData
	dolphin: AnimalData
	zebra: AnimalData

This would work but it may lead to bugs later if someone added a new field to AnimalSpecies without adding a corresponding field to AnimalDict. I wanted to find an easier way to do this.

Final Solution

After digging through the documentation, I discovered that I could achieve my goals using custom root types.

Using the RootModel class, we can rewrite the AnimalDict class as follows:

from pydantic import RootModel
from typing import Dict

class AnimalDict(RootModel):
    root: Dict[AnimalSpecies, AnimalData]

    def __getitem__(self, key: AnimalSpecies):
        return self.root[key]

    def __setitem__(self, key: AnimalSpecies, value: AnimalData):
        self.root[key] = value

    @model_validator(mode="after")
    def check_dictionary_types(self) -> Self:
        keys = self.root.keys()
        for key in keys:
            if not isinstance(key, AnimalSpecies):
                raise TypeError(f"Dictionary key {key} isn't of type AnimalSpecies")

            value = self.root[key]
            if not isinstance(value, AnimalData):
                raise TypeError(
                    f"The object {value} bound to key {key} isn't of type AnimalData"
                )

        return self

Here’s how to use this class in practice:

animals = AnimalDict({})
animals[AnimalSpecies.LION] = AnimalData(
    habitat="cave", diet="meat")

The model validator function check_dictionary_types allows the model_validate method to ensure that the root object obeys the schema. This function will trigger a type error if either a key or a value isn’t of the appropriate type.

Here’s an example:

example = AnimalDict({})
example["o"] = {"b": 2}
# ⚠️ This line below will trigger an error ⚠️
AnimalDict.model_validate(example)

Conclusion

Pydantic’s custom root types provide a flexible solution to define and validate complex types that might be challenging to express using other methods. By leveraging the RootModel class, I created a dictionary type for my specific use case.

This approach provides a scalable way to validate custom dictionary types throughout my codebase.