How does Pydantic discriminator works?

El discriminador de Pydantic permite la definición de estructuras de datos con múltiples tipos, utilizando un campo de discriminador para determinar el tipo real del objeto. Esto habilita la validación de tipos y la serialización/deserialización basada en el valor del discriminador, asegurando la integridad de los datos y la flexibilidad en la representación de diferentes tipos de objetos.

A partir de Pydantic 1.9, lo tenemos disponible. Let’s showcase it in an easy way:

from pydantic import BaseModel, Field, parse_obj_as
from typing import Literal, Union, Annotated

class Tiger(BaseModel):
    animal_type: Literal["tiger"] = "tiger"
    ferocity_scale: float = Field(..., ge=0, le=10)

class Shark(BaseModel):
    animal_type: Literal["shark"] = "shark"
    ferocity_scale: float = Field(..., ge=0, le=10)

class Lion(BaseModel):
    animal_type: Literal["lion"] = "lion"
    ferocity_scale: float

class WildAnimal(BaseModel):
    __root__: Annotated[Union[Tiger, Shark, Lion], Field(..., discriminator='animal_type')]

my_shark = WildAnimal.parse_obj({'animal_type': 'shark', 'ferocity_scale': 5}).__root__
#print(Shark(ferocity_scale=5).json())

# Desarialice
WildAnimal.parse_raw(Shark(ferocity_scale=5).json())

## WildAnimal(__root__=Shark(animal_type='shark', ferocity_scale=5.0))

print(isinstance(my_shark, Shark))

## True

Puedes encontrar este ejemplo de polimosfirmo junto con algun ejemplo de código, además de alguna interesante discusión en: https://github.com/pydantic/pydantic/discussions/5785

Ejemplo de discriminador de unión anotada Pydantic

Pero podríamos utilizar un enfoque muy simple para lograr la mayoría de los usos mediante el uso de la unión Annotated.

Animal = Annotated[Union[Tiger, Shark], Field(discriminator='animal_type')]
raw_data = {
    "animal_type": "tiger",
    "ferocity_scale": 6
}
parse_obj_as(Animal, raw_data)

## Tiger(animal_type='tiger', ferocity_scale=6.0)

Prepárate para la magia de la clase Field, cortesía de Pydantic. Está equipada con un poder especial llamado “discriminator”. Al configurar el discriminador en “pet_type”, desbloqueamos la capacidad de distinguir entre nuestras criaturas fantásticas. ¡Es como darles su propio foco especial!

¡Agárrate fuerte, porque estamos a punto de adentrarnos en las tierras salvajes de raw_data. Guarda los secretos de un “pet_type” con el espíritu ardiente de un “tigre” y un fascinante recuento de “rayas” de 6. Es como si estuviéramos mirando a un zoológico digital.

¡Y ahora, es hora del espectáculo! Invocamos al poderoso parse_obj_as para que haga su magia de codificación. Le presentamos a nuestro majestuoso Animal y al enigmático raw_data. ¡Abra Kadabra! Con un movimiento de su varita, la transformación se despliega. Los datos crudos se convierten en una impresionante representación de nuestro Animal elegido. ¡Es como una metamorfosis mágica!

Example of Polimorfic Base Model

class PolymorphicBaseModel(BaseModel):
    type: str

    _subtypes = dict()

    def __init_subclass__(subcls, type=None, **kwargs):
        super().__init_subclass__(**kwargs)
        if type:
            # n.b. if a subclass declares its own _subtypes dict, it'll take precedence over this one.
            # This would allow us to re-use the same type names across different classes.
            if type in subcls._subtypes:
                raise AttributeError(
                    f"Class {subcls} cannot be registered with polymorphic type='{type}' because it's already registered "
                    f" to {subcls._subtypes[type]}"
                )
            subcls._subtypes[type] = subcls
    @classmethod
    def _convert_to_real_type(cls, data):
        data_type = data.get("type")

        if data_type is None:
            raise ValueError(f"Missing 'type' for {cls}")

        subcls = cls._subtypes.get(data_type)

        if subcls is None:
            raise TypeError(f"Unsupported sub-type: {data_type}")
        if not issubclass(subcls, cls):
            raise TypeError(f"Inferred class {subcls} is not a subclass of {cls}")

        return subcls(**data)

    @classmethod
    def parse_obj(cls, data):
        return cls._convert_to_real_type(data)
    
    
class Animal(PolymorphicBaseModel):
    name: str
    color: str = None

class Cat(Animal, type="cat"):
    type: Literal["cat"] = "cat"
    hairless: bool

class Dog(Animal, type="dog"):
    type: Literal["dog"] = "dog"
    breed: str

cat_instance = Animal.parse_obj({"type":"cat", "hairless": False, "name": "meaw", "color": "black"})
print(isinstance(cat_instance, Cat))

## True

El PolymorphicBaseModel, una clase base que sienta las bases para el comportamiento polimórfico. Define un atributo de tipo requerido e introduce un diccionario oculto _subtypes para llevar un registro de los subtipos.

A continuación, nos sumergimos en el método init_subclass, donde sucede la magia. Permite que las subclases se registren a sí mismas con un tipo polimórfico específico. Esto nos permite distinguir entre diferentes subtipos dentro de la jerarquía de PolymorphicBaseModel.

¡Pero espera, hay más por descubrir! Hacemos uso del método _convert_to_real_type, encargado de convertir los datos a su subtipo real en función del atributo de tipo proporcionado. Comprueba si el tipo es válido, encuentra la subclase correspondiente y asegura que sea una subclase válida de la clase base.

Finalmente, llegamos al método parse_obj, donde tiene lugar el verdadero análisis. Sirve como punto de entrada para analizar objetos de la jerarquía polimórfica. Utilizando el método _convert_to_real_type, transforma los datos en una instancia de la subclase adecuada.

¡Y ahí lo tienes! Un vistazo al mundo de los modelos polimórficos. Es un mundo donde las clases base y los subtipos se unen, permitiendo un análisis de objetos flexible y dinámico. Aprovecha el poder del polimorfismo y permite que tu código se adapte y evolucione con elegancia.

Pydantic 2: TypeAdapter para analizar datos en una unión discriminada

En Pydantic v2, puedes utilizar el TypeAdapter para analizar datos en una unión discriminada. Sin embargo, ten en cuenta que Pydantic v2 se encuentra actualmente (2023-06-18) en prelanzamiento, y la versión actual del módulo es la v1.7.

Por lo tanto, asegúrate de actualizar a Pydantic v2 cuando esté disponible para aprovechar esta característica.

from pydantic import TypeAdapter

adapter = TypeAdapter(Annotated[Union[Child1, Child2], Field(discriminator='type')])

child = adapter.validate_json(my_json_data)

Mantente al tanto de consejos sobre Pydantic y Python

Esperamos que esta publicación te haya ayudado a familiarizarte con el uso de uniones y discriminadores en Pydantic, mostrando algunas de sus funcionalidades y permitiéndote disfrutar de sus beneficios.

Si deseas mantenerte actualizado…