Pydantic, Enums e IntEnums. Una historia de validación

Validación y tipado de Datos en Python: Dominando Pydantic, Enums e IntEnums para Aplicaciones Python Robustas.

Logo del módulo de Python Pydantic.

¿Qué es un Enum y su integración con Pydantic?

¡Oh, esta es una gran pregunta! ¡Nunca la había escuchado antes! Un Enum, para los no iniciados, es una característica ingeniosa que, cuando se combina con la biblioteca Pydantic, te ayuda a controlar el caos de la jungla de datos. Es una combinación del Enum de Python (abreviatura de enumeración) y el poder de validación de Pydantic. Los Enums te permiten definir un conjunto de valores con nombres a los que tus datos deben adherirse. Pydantic luego verifica si tus datos son parte de este club exclusivo y, si no lo son, amablemente les muestra la puerta de salida.


¿Cómo me ayudan Pydantic y Enum?

Si alguna vez has trabajado con datos, sabes que tienen voluntad propia. A veces son impecables y perfectos. En otras ocasiones, son un desastre total. Los Enums están aquí para salvarte de esos días en los que tus datos deciden actuar por su cuenta.

Al definir Enums y utilizarlos en tus modelos de Pydantic, creas un conjunto de reglas que tus datos deben seguir, como un profesor estricto pero justo. De esta manera, puedes asegurarte de que solo datos válidos entren en tu sistema y que cualquier valor atípico se maneje adecuadamente.

from enum import Enum
from pydantic import BaseModel, ValidationError

class Pet(BaseModel):
  name: str
  animal_type: str
  sex: str
  

Puedes ser más específico al definir tus modelos de Pydantic utilizando Enums:

from pydantic import ValidationError

class Sex(Enum):
  MALE = 'male'
  FEMALE = 'female'
  
class DomesticAnimals(Enum):
  CAT = 'cat'
  DOG = 'dog'
  FISH = 'fish'
  BIRD = 'bird'

# And then

class Pet(BaseModel):
  name: str
  animal_type: DomesticAnimals
  sex: Sex
  
Pet(name='Timmy', animal_type='bird', sex='male')
## Pet(name='Timmy', animal_type=<DomesticAnimals.BIRD: 'bird'>, sex=<Sex.MALE: 'male'>)

Agregar un valor que no existe (por ejemplo, “Tigre”, que no está permitido en animales domésticos) a un Enum existente generará un error el valor no es un miembro de la enumeración válido. Este es uno de los principales casos de uso para ellos.

import pytest

with pytest.raises(ValidationError, match=' value is not a valid enumeration member') as e_info:
  Pet(
    name='Timmy', 
    animal_type='tiger', 
    sex='male'
  )
print(e_info.value)
## 1 validation error for Pet
## animal_type
##   value is not a valid enumeration member; permitted: 'cat', 'dog', 'fish', 'bird' (type=type_error.enum; enum_values=[<DomesticAnimals.CAT: 'cat'>, <DomesticAnimals.DOG: 'dog'>, <DomesticAnimals.FISH: 'fish'>, <DomesticAnimals.BIRD: 'bird'>])

(*Note the UPPER_CASE_NOTATION)


¿Y qué hay de IntEnum? ¿Cuál es la diferencia entre Enum e IntEnum?

En resumen, la principal diferencia entre Enum e IntEnum radica en el tipo de valores que representan. Enum es una clase de enumeración genérica que puede utilizarse con cualquier tipo de datos, mientras que IntEnum está específicamente diseñado para valores enteros y permite la comparación directa con enteros.

Existen dos beneficios principales al utilizar IntEnums en los casos de uso correctos:

  • Dado que IntEnum garantiza que todos los miembros de la enumeración tienen un valor entero, es posible ordenarlos.
  • Los miembros de IntEnum pueden compararse directamente con enteros, mientras que los miembros de Enum no pueden utilizarse en operadores de comparación de enteros.
from enum import IntEnum
import pytest

class ResponseCode(IntEnum):
    OK = 200
    NOT_FOUND = 404
    ERROR = 500
    
assert ResponseCode.OK == 200
assert ResponseCode.OK <= ResponseCode.NOT_FOUND

with pytest.raises(TypeError, match='cannot extend enumeration') as e_info: # Check that a TypeError is raised
  class ExtendedResponseCode(ResponseCode):
      CUSTOM = 300
print(e_info.value)
## ExtendedResponseCode: cannot extend enumeration 'ResponseCode'


¿Es posible crear una subclase de un Enum (o un StrEnum / IntEnum)?

No se supone que sea posible. Si intentas hacerlo mediante la herencia directa, es posible que se genere un TypeError. Para explicarlo, la documentación dice: “Permitir la creación de subclases de enums que definen miembros llevaría a una violación de algunas invariantes importantes de tipos e instancias.”

¿Qué estaríamos infringiendo? Veamos el comentario de Guido en (2013):

from enum import Enum

class Color(Enum):
  red = 1
  green = 2
  blue = 3

class MoreColor(Color): # this is not possible as we've seen
  cyan = 4
  magenta = 5
  yellow = 6

type(MoreColor.red) is Color

type(MoreColor.red) is not MoreColor

#En otras palabras, mientras 'red' es accesible en MoreColor
#es realmente una instancia de Color?

#Vaya, esto es un caos. No queremos que MoreColor.red y
#Color.red sean objetos diferentes, pero usando isinstance() check
#parece confuso.

#not isinstance(Color.red, MoreColor)
#isinstance(MoreColor.yellow, Color)

En algunas versiones de Python, esto funciona sin mostrar un mensaje de error, pero es un comportamiento no deseado.

class Color(Enum):
  red = 1
  green = 2
  blue = 3

class MoreColor(Enum, Color):
  cyan = 4
  magenta = 5
  yellow = 6

Uno podría argumentar que las enumeraciones existen para garantizar la exclusión mutua sobre un conjunto finito no ordenado. Agregar miembros adicionales a una enumeración existente no viola esta garantía. Por lo tanto, si estás seguro de tu caso de uso y de lo que estás haciendo, es posible crear una solución alternativa. Una solución limpia utilizando un decorador es:

from enum import Enum
from typing import Any, Callable

class EnumBase(Enum):
    def __eq__(self, other: Any) -> bool:
        if isinstance(other, Enum):
            return self.value == other.value
        return False
      
def extend_enum(parent_enum: EnumBase) -> Callable[[EnumBase], EnumBase]:
    """Decorator function that extends an enum class with values from another enum class."""
    def wrapper(extended_enum: EnumBase) -> EnumBase:
        joined = {}
        for item in parent_enum:
            joined[item.name] = item.value
        for item in extended_enum:
            joined[item.name] = item.value
        return EnumBase(extended_enum.__name__, joined)
    return wrapper
class Parent(EnumBase):
  A = 1
  B = 2
  
@extend_enum(Parent)
class ExtendedParent(EnumBase):
  C = 3
  
print(
type(Parent.A) is Parent,
type(Parent.A) is not ExtendedParent,
Parent.A == ExtendedParent.A
)
## True True True

Pero esta no es una solución perfecta, ya que tiene algunas desventajas o limitaciones de las que debes ser consciente. En este caso, un Enum no relacionado (llamado RandomEnum) que implementa el mismo valor de enumeración es igual en la comparación con nuestras clases Parent y ExtendedParent:


class RandomEnum(EnumBase):
  A = 1
  
Parent.A == RandomEnum.A == ExtendedParent.A
## True


Mantente actualizado en consejos de Pydantic y Python

Esperamos que esta publicación te haya ayudado a familiarizarte con el uso de Enum en Pydantic y te haya permitido disfrutar de una presentación de algunas de sus funcionalidades.

Si deseas mantenerte actualizado…

Python Pydantic
Carlos Vecina
Carlos Vecina
Senior Data Scientist at Jobandtalent

Senior Data Scientist at Jobandtalent | AI & Data Science para aportar valor en la empresa