Crear y manejar enumerados (enums) en Python

Tiempo de lectura: 2 minutos

Hoy vamos a aprender cómo podemos crear y manejar datos enums de enumerado con Python.

Los enums nos sirven para organizar los índices dentro de nuestro código. En el que podemos incluir los roles, los tipos de algun aviso / alerta, colores etc…

Para crear esta estructura tenemos que poner lo siguiente:

class Tipo_aviso(Enum):
    ALERTA = 1
    MODERADO = 2
    INFORMACION = 3

En este caso he creado 3 tipos de aviso, que se corresponden a la tabla de la base de datos que contiene este índice. Los he asignado con número, también pueden ser texto en forma de clave.

Ahora podemos asignarlo a un objeto:

class Aviso(BaseModel):
    id: Optional[int] = None
    id_usuario_destino: int
    titulo: str
    descripcion: str
    tipo_aviso: Tipo_aviso
    timestamp: Optional[datetime.datetime] = None

Recuerda que en Python, primero tienes que añadir el objeto que luego vas a asignar ya que sigue un orden.

Ahora podremos crear el objeto de esta forma:

aviso: Aviso = Aviso(
        id_usuario_destino=1,
        titulo="Prueba",
        descripcion="Prueba descripción",
        tipo_aviso=Tipo_aviso.ALERTA
    )

Y nos asignará el valor de Alerta.

Para obtener los valores del enum hacemos lo siguiente:

print(aviso.tipo_aviso)             # Tipo_aviso.ALERTA
print(aviso.tipo_aviso.name)        # 'ALERTA'
print(aviso.tipo_aviso.value)       # 1

Como has visto, el enum devuelve una estructura clave/valor. Para obtener el valor tendremos que indicar .value y para obtener la clave tendremos que indicar .name

A veces nos interesa que el propio objeto que implementa este enum nos devuelva directamente el value. Para ello añadimos lo siguiente:

class Config:
        use_enum_values = True

Con esta configuracion personalizada del objeto indicamos que obtenga directamente el valor del enum (entero).

El objeto completo queda así:

class Aviso(BaseModel):
    id: Optional[int] = None
    id_usuario_destino: int
    titulo: str
    descripcion: str
    tipo_aviso: Tipo_aviso
    timestamp: Optional[datetime.datetime] = None

    class Config:
        use_enum_values = True

De esta forma evitamos problemas al enviar el objeto mediante un JSON o al insertarlo en la Base de datos, además de al recuperarlo.