Hace tiempo hablé en otra entrada
sobre las ventajas de evitar el uso de primitivos en nuestro código. En aquel momento propuse la
creación de clases hoy en cambio vamos a ver otra alternativa, los Enums.
El concepto Enum en Python
A diferencia de otros lenguajes, en Python no existe como tal el concepto de Enum tradicional.
En Python un Enum no es ni más ni menos que una clase que permite definir una serie de constantes a las que se les asigna un nombre y un valor. Pero como clases que son nuestros enums serán capaces de hacer muchas cosas más y que son muy habituales en orientación a objetos.
Algunos de los conceptos clave para entender cómo trabajar con enums en Python son:
- Para crear un enum necesitamos crear una clase que herede de la clase
Enum. - Los miembros de un enum tienen tanto nombre como valor.
- Por defecto los enums permiten que diferentes nombres dentro de un enum tengan el mismo valor.
- Para evitar este comportamiento debemos usar el decorador
@unique
- Para evitar este comportamiento debemos usar el decorador
- Los enums en el fondo son clases por lo que podrán tener métodos y propiedades como cualquier otra.
Caso práctico usando primitivos
Después de toda esta teoría vamos a ver un caso práctico real (muy simplificado) donde en función del tipo de ingeniero dentro del sistema debemos calcular su coste total para la empresa.
Por simplicidad vamos a asumir que la propiedad salary es un float, pero está mal hecho,
debería ser algo tipo Amount o similar.
Como siempre vamos a empezar por los tests:
class TestEnums:
@pytest.mark.parametrize(
"engineer_type,salary,expected_total_cost",
[
("Junior", 100000, 100000),
("Senior", 5000, 7500),
("Staff", 3000, 6000),
("Team Lead", 1000, 3000),
("Enginering Manager", 1, 4),
],
)
def test_enums(self, engineer_type: str, salary: float, expected_total_cost: int):
engineer = Engineer(name="Ada Lovelace", type=engineer_type, salary=salary)
total_cost = calculate_engineer_cost(engineer)
expect(total_cost).to(equal(expected_total_cost))
Una vez tenemos nuestro test vamos a ir creando el mínimo código posible para ir eliminando los errores y que nuestro código al menos sea ejecutable.
Para ello creamos la clase Engineer:
@dataclass
class Engineer:
name: str
salary: float
type: str
Lo siguiente que necesitamos implementar es nuestra lógica de negocio.
def calculate_engineer_cost(engineer: Engineer) -> float:
return engineer.salary * hour_rate(engineer.type)
def hour_rate(engineer_type: str) -> float:
return {
"Junior": 1,
"Senior": 1.5,
"Staff": 2,
"Team Lead": 3,
"Engineering Manager": 4,
}[engineer_type]
En este punto ya estamos en condición de ejecutar los tests sin errores. Desafortunadamente no todo son buenas noticias_:
KeyError: 'Enginering Manager'
¿Qué está ocurriendo? ¿De dónde proviene esa excepción de tipo KeyError?
Efectivamente, el problema viene de que estamos accediendo con una key inválida (nuestro engineer_type)
en nuestro método hour_rate. ¿Pero cómo es posible? Si aún no lo ves, por favor vuelve a mirar los parámetros que introducimos a nuestro test y mira detenidamente el último ejemplo del @pytest.mark.parametrize.
Efectivamente hay un error en Enginering Manager.
Es Engineering, NO Enginering.
Este error puede parecer una tontería, pero es más habitual de lo que parece y además es el típico error que o bien lo ves de manera casi instantánea o puedes estar una hora dando vueltas sin fijarte.
El origen de este error es que estamos usando el tipo string para definir nuestro engineer_type lo
que permite que cualquier string pueda ser pasado al método hour_rate. Además tenemos por diferentes lugares de nuestra aplicación los mismos strings repetidos una y otra vez, en lugar de tener un único sitio donde se define su valor correcto y el resto de la aplicación hace
uso del mismo.
Caso práctico usando Enums
Vamos a resolver el mismo problema, pero esta vez haciendo uso de un enum en lugar de un string.
Primero actualizamos los tests acordes a como esperamos que vaya a ser definido nuestro nuevo enum:
class TestEnums:
@pytest.mark.parametrize(
"engineer_type,salary,expected_total_cost",
[
(EngineerType.JUNIOR, 100000, 100000),
(EngineerType.SENIOR, 5000, 7500),
(EngineerType.STAFF, 3000, 6000),
(EngineerType.TEAM_LEAD, 1000, 3000),
(EngineerType.ENGINEERING_MANAGER, 1, 4),
],
)
def test_enums(self, engineer_type: EngineerType, salary: int, expected_total_cost: int):
engineer = Engineer(name="Ada Lovelace", type=engineer_type, salary=salary)
total_cost = calculate_engineer_cost(engineer)
expect(total_cost).to(equal(expected_total_cost))
Aquí ya podemos ver como desaparecen los magic strings y pasamos a usar directamente el enum que además tiene autocompletado por parte del IDE por lo que es casi imposible confundirse.
Debemos por tanto crear un enum llamado EngineerType que tendrá 5 posibles valores:
@unique
class EngineerType(Enum):
JUNIOR = "Junior"
SENIOR = "Senior"
STAFF = "Staff"
TEAM_LEAD = "Team Lead"
ENGINEERING_MANAGER = "Engineering Manager"
Una vez creado el enum vamos a actualizar nuestra clase Engineer para que el tipo deje de ser un string:
@dataclass
class Engineer:
name: str
salary: float
type: EngineerType
Por último solo nos queda adaptar nuestro método hour_rate para que el parámetro que recibe sea del tipo EngineerType:
def calculate_engineer_cost(engineer: Engineer) -> float:
return engineer.salary * hour_rate(engineer.type)
def hour_rate(engineer_type: EngineerType) -> float:
return {
EngineerType.JUNIOR: 1,
EngineerType.SENIOR: 1.5,
EngineerType.STAFF: 2,
EngineerType.TEAM_LEAD: 3,
EngineerType.ENGINEERING_MANAGER: 4,
}[engineer_type]
Si ahora ejecutamos los tests todo debería estar correcto.
Personalmente creo que nuestro código es mucho más robusto ahora y el esfuerzo para conseguirlo ha sido relativamente bajo.
Los posibles valores de nuestros engineer_type solo se definen en un sitio, no es posible utilizar otros tipos
que no están definidos, nuestro código es mucho más semántico y hemos reducido la carga cognitiva
al no tener que saber cuál es el valor del tipo ENGINEERING_MANAGER.
Los Enums en Python son mucho más que enums
Como decía al principio, en Python los enums no dejan de ser clases por lo que podemos añadirles lógica si lo creemos necesario.
Imaginemos por ejemplo que queremos proporcionar un constructor a nuestros usuarios
cuando quieran crear un EngineerType de forma dinámica. Es posible que alguien trate
de crear un EngineerType con un valor incorrecto y en ese caso queremos indicarle cuáles
son las opciones posibles.
Una vez más empezamos por el test:
class TestEnums:
...
def test_raise_an_error_when_creating_with_an_invalid_type(self) -> None:
invalid_type = "invalid-type"
error_message = f"Sorry the type '{invalid_type}' is not correct, please use one of {EngineerType.all_values()}"
try:
EngineerType.from_(invalid_type)
except InvalidEngineeringTypeException as ex:
expect(str(ex)).to(equal(error_message))
Vamos a crear el método from_(value) en nuestro EngineerType:
@unique
class EngineerType(Enum):
...
@staticmethod
def all_values() -> list[str]:
return [item.value for item in list(EngineerType)]
@staticmethod
def from_(value: str) -> "EngineerType":
try:
return EngineerType(value)
except ValueError as ex:
raise InvalidEngineeringTypeException(value) from ex
Y por último nuestra excepción InvalidEngineeringTypeException:
class InvalidEngineeringTypeException(Exception):
def __init__(self, invalid_type: str) -> None:
error_message = f"Sorry the type '{invalid_type}' is not correct, please use one of {EngineerType.all_values()}"
super().__init__(error_message)
Si volvemos a lanzar los tests todo debería estar en verde!
Conclusiones
La idea detrás de esta entrada no es más que dar a conocer otra herramienta más que podamos incluir en nuestra caja de herramientas cuando estamos programando.
Que sepamos que existen los Enums y que en ciertas ocasiones nos pueden venir muy bien. Para mi son un compañero ideal cuando tengo que definir una entidad con un número fijo y limitado de opciones.
Nos vemos en la próxima, un saludo!
El codigo de todos estos ejemplos esta disponible AQUÍ por si quieres echarle un vistazo con calma.