#!/usr/bin/env python
# coding: utf-8

# # Using Pydantic with custom classes
# 
# [Pydantic](https://docs.pydantic.dev/latest/) is an awesome data validation library for Python.
# It’s getting very popular these days because it’s fast, has lots of features and is pleasant to use.
# I’m using it at work and it personal projects.
# It’s one of the strong points of [FastAPI](https://fastapi.tiangolo.com), the new popular Python framework.
# It can even integrate with Django thanks to [django-ninja](https://django-ninja.dev/.
# 
# While it integrates with standard Python types out of the box, using it with custom classes requires a bit more work.
# [It’s documented](https://docs.pydantic.dev/latest/concepts/types/#custom-types), but I don’t think this part of the documentation is very clear.
# That’s why, in this article, I’ll share my experience with how I integrated a custom class with Pydantic.
# I hope you can use it to understand better how it works and thus facilitate using your custom classes with Pydantic.
# In a second article (coming soon), I’ll give another example of custom classes with Pydantic by using enums whose members are Pydantic models.
# 
# Feel free to download these files and experiment!
# The code is tested on Python 3.11 but should work with any recent enough version of Python.
# I suppose you know Python and have a basic understanding of Pydantic, generic types in Python and know the common dunder methods.
# 
# .. contents:: Table of contents
# 
# 
# ## Why I needed to integrate a custom class in the first place
# 
# My goal was to create a read only mapping compatible with Pydantic.
# Here’s how it must work:
# 
# - Usage with Pydantic must be transparent:
#   - When the model is serialized, the mapping must be serialized as a dict.
#   - It must be created from a dict containing values.
#     Both the keys and the values from this input dict must be parsed and validated.
#     To increase reusability, the model must be made generic to allow validation of different types: for instance a `ReadonlyMapping[str, int]` and `ReadonlyMapping[int, float]`.
#   - Just like any other Pydantic models, if I pass an instance to the model, the model must use the class directly.
# - Once set in the model, the field must not be editable any more.
#   The goal being that the model must be fully immutable so that neither its fields nor the content of the mapping can be edited by mistake.
# 
# So my goal is to arrive at something like this:

# In[1]:


from pydantic import BaseModel, ConfigDict, ValidationError


class ImmutableModel(BaseModel):
    # Make the model immutable: cannot add nor change attributes.
    model_config = ConfigDict(frozen=True)

    my_field: int
    my_dict: dict[str, str]


m = ImmutableModel.model_validate({"my_field": 1, "my_dict": {}})

try:
    m.my_field = 2
    assert False, "Must fail"
except ValidationError as e:
    print("Fails as expected with:", e)

# The underlying structure is still mutable and can be changed.
# despite the model being frozen. That’s what we want to prevent.
m.my_dict["key"] = "value"


# ## The class without Pydantic
# 
# Let’s start by building a class that will be our generic readonly mapping.
# At its core, this class will just be a wrapper around a good old dict.
# It will rely on `typing.Generic` to be generic and will implement `__getitem__` to provide item access.
# It’ll also implement `__str__` and `__repr__` to ease debugging.
# `__setitem__` won’t be implemented to prevent setting an item and get a `TypeError` if we try to.
# 
# .. note:: I didn’t use `MappingProxyType` because I wanted a class that’s easier to extend to fit my needs here. In another context, that’s the solution I’d have chosen to have a readonly mapping.

# In[2]:


from typing import TypeVar, Generic


# Generic type for the key.
Key = TypeVar("Key")
# Generic type for the value.
Value = TypeVar("Value")


class ReadonlyMapping(Generic[Key, Value]):
    def __init__(self, mapping: dict[Key, Value] | None = None):
        self._mapping: dict[Key, Value] = mapping or {}

    def __getitem__(self, key: Key):
        return self._mapping[key]

    def __str__(self):
        return str(self._mapping)

    def __repr__(self):
        return repr(self._mapping)


# This class can then be used this way:

# In[3]:


readonly_mapping = ReadonlyMapping[int, str]({1: "Test"})
print(readonly_mapping[1])


# And it fails as expected on key assignment:

# In[4]:


try:
    readonly_mapping[2] = "Testing"
    assert False, "Must fail"
except TypeError as e:
    print("Fails as expected with:", e)


# ## Adding Pydantic
# 
# Now that the base is built, let’s iterate on that.
# To hook Pydantic with this custom class, a class method named `__get_pydantic_core_schema__` must be added.
# It’ll use low level Pydantic functions to generate a proper Pydantic schema based on your generic arguments and let Pydantic handle the parsing and the validation.
# These methods come from `pydantic_core` and are used internally by Pydantic.
# They are documented [here](https://docs.pydantic.dev/latest/api/pydantic_core_schema/).
# Here’s how it looks like (with inline comments to explain its tricky part):

# In[5]:


from pydantic import BaseModel, GetCoreSchemaHandler, ValidationError
from typing import Any, Generic, TypeVar
from pydantic_core import core_schema
from typing import get_args as get_typing_args


Key = TypeVar("Key")
Value = TypeVar("Value")


class ReadonlyMapping(Generic[Key, Value]):
    def __init__(self, mapping=None):
        self._mapping: dict[Key, Value] = mapping or {}

    def __getitem__(self, key: Key):
        return self._mapping[key]

    def __str__(self):
        return str(self._mapping)

    def __repr__(self):
        return repr(self._mapping)

    @classmethod
    def __get_pydantic_core_schema__(
        cls, source: Any, handler: GetCoreSchemaHandler
    ) -> core_schema.CoreSchema:
        # Retrieve the generic arguments. args is a tuple of the type used for
        # Key and the type used for Value.
        args = get_typing_args(source)
        # If no generics are passed, args will be an empty tuple.
        # Let’s fallback to a raw dict in this case.
        if args:
            # Make Pydantic generate a schema for dict[Key, Value] that it can
            # directly use for validation.
            # This allows for Pydantic to parse the keys and the values and to
            # validate them without further work.
            validate_dict_data_schema = handler.generate_schema(
                dict[args[0], args[1]]
            )
        else:
            validate_dict_data_schema = handler.generate_schema(dict)

        # Create a schema that will build our instance from cls
        # (ie ReadonlyMapping) from the validated dict.
        from_dict_schema = core_schema.no_info_after_validator_function(
            cls, validate_dict_data_schema
        )

        # Validator to hook an already built instance of ReadonlyMapping
        # into the model.
        from_instance_schema = core_schema.is_instance_schema(cls)

        # Combine from_dict_schema with from_instance_schema to be able
        # to use both.
        # They will be tried in order: 1st from_instance_schema and if
        # it fails from_dict_schema.
        # If the last one fails, validation will fail.
        schema = core_schema.union_schema(
            [from_instance_schema, from_dict_schema]
        )

        return core_schema.json_or_python_schema(
            # Use our schema when raw JSON data is used directly.
            json_schema=schema,
            # With Python data.
            python_schema=schema,
            # When the model is serialized to Python or JSON, return the
            # underlying dict that Pydantic can handle directly.
            serialization=core_schema.plain_serializer_function_ser_schema(
                lambda instance: instance._mapping
            ),
        )


# ### Testing
# 
# Let’s add a test model for testing:

# In[6]:


class MyModel(BaseModel):
    my_mapping: ReadonlyMapping[str, int]


# And on to the tests we go!
# 
# #### Loading from a `dict`

# In[7]:


m = MyModel.model_validate({"my_int": 1, "my_mapping": {"key": "77"}})
print("The model:", type(m), m)
m_python = m.model_dump(mode="python")
print("As Python:", type(m_python), m_python)
m_json = m.model_dump_json()
print("As JSON:", type(m_json), m_json)


# #### Loading from json

# In[8]:


m = MyModel.model_validate_json("""{"my_mapping": {"key": "77"}}""")
print("The model from JSON:", m)


# #### Key access

# In[9]:


print(m.my_mapping, m.my_mapping["key"])


# #### Access invalid key

# In[10]:


try:
    print(m.my_mapping["nope"])
    assert False, "Must fail"
except KeyError as e:
    print("Accessing an invalid key fails as expected with error:", e)


# #### Direct creation

# In[11]:


m = MyModel(my_mapping=ReadonlyMapping({"key": 77}))
print("Direct creation:", m)


# #### Validating with a ReadonlyMapping instance

# In[12]:


m = MyModel.model_validate({"my_mapping": ReadonlyMapping({"key": 77})})
print("Validating with a ReadonlyMapping instance:", m)


# #### Setting a key

# In[13]:


try:
    m.my_mapping["key"] = 0
    assert False, "Must fail"
except TypeError as e:
    print("Setting a key fails as expected with error:", e)


# #### Key parsing

# In[14]:


class MyOtherModel(BaseModel):
    mapping: ReadonlyMapping[int, int | bool]


m = MyOtherModel.model_validate({"mapping": {"1": "1", "2": False}})
print("Keys and values are correctly parsed and validated:", m)


# #### Validation failure

# In[15]:


try:
    MyOtherModel.model_validate({"mapping": {"key": "1", "2": "Some value"}})
    assert False, "Must fail"
except ValidationError as e:
    print("This fails as expected with error e:", e)


# #### Without any generic annotation:

# In[16]:


class NoTypeAnnotationModel(BaseModel):
    mapping: ReadonlyMapping

m = NoTypeAnnotationModel.model_validate({"mapping": {"key": "value", 1: 1}})
print("Read and parsed, no validation:", m)


# ## Conclusion
# 
# I hope I demystified using a custom class with Pydantic.
# Once you know a bit how `__get_pydantic_core_schema__` can be used and how to use the core validation schemas, it’s not that hard.
# Just like me, you’ll probably have to experiment a bit to fully grasp what is going on.
# And if something is unclear, don’t hesitate to ask a question below!