Dynamic generation of informative `Enum`s in Python

Ignacio Vergara Kausel published on

3 min, 444 words

Categories: dev

Since I started learning some Rust and getting to know its enums, I've been very into using Enums in Python. And while enums in both languages are very different, they're close enough for most of my current use cases. Moreover, within Python, Enums are used in frameworks like Typer (CLI) and FastAPI (API), where they are used to provide validation to input within a set of possibilities.

In Python, an Enum without explicit member values is created as follows

from enum import Enum, auto
class Color(Enum):
    RED = auto()
    BLUE = auto()
    GREEN = auto()

list(Color)
# [<Color.RED: 1>, <Color.BLUE: 2>, <Color.GREEN: 3>]

or alternatively, using a functional API

from enum import Enum
Animal = Enum('Animal', 'ANT BEE CAT DOG', module=__name__)
list(Animal)
# [<Animal.ANT: 1>, <Animal.BEE: 2>, <Animal.CAT: 3>, <Animal.DOG: 4>]

Great! The problem is that when having Enums that interface with external users, the integer values of the members are of little use and very uninformative.

That's why I tend to prefer to use a different way, where the values are strings.

class Animal(str, Enum):
    ANT = 'ant'
    BEE = 'bee'
    CAT = 'cat'
    DOG = 'dog'
list(Animal)
# [<Animal.ANT: 'ant'>,
#  <Animal.BEE: 'bee'>,
#  <Animal.CAT: 'cat'>,
#  <Animal.DOG: 'dog'>]

Now, attending to the real problem at hand hinted at by the title. I had a rather long set of about 15 possible elements, and I didn't want to manually create all the elements of the Enum myself in a way that the values are informative. That's the exact thing we try to avoid, repetitive manual work. We already saw the functional API work with an iterable, so basically, the problem is already solved. However, it does in a way that the member values are uninformative.

The solution is to provide a dictionary instead of a list (or a string in the example). Then, the (keys, value) pairs of the dictionary become the member name and value of the Enum as shown below. Going from a list (or an iterable) to a dictionary is straightforwardly achieved with a dictionary comprehension.

Animal = Enum('Animal', {el:el.lower() for el in 'ANT BEE CAT DOG'.split(" ")}, type=str)
list(Animal)
# [<Animal.ANT: 'ant'>,
#  <Animal.BEE: 'bee'>,
#  <Animal.CAT: 'cat'>,
#  <Animal.DOG: 'dog'>]

The ultimate test, to check wether both implementations are the same, would be to check the MRO of each of the classes generated. Doing so via Animal.mro() for each one, the same output is obtained, namely, [<enum 'Animal'>, str, <enum 'Enum'>, object].

Ultimately, the idea and solution are not too difficult, but was a nice exploration that accomplishes what I was going for.

Share: