How to use strings as name aliases in Python enums

Python has support for enumerations built into the standard library since version 3.4. The Enum type is quite powerful but serializing enum members to human readable representations and deserializing them to an enum meber can be cumbersome.

If we want to have an enumeration of certain train stations between the two Austrian cities Wien and Wels, the following approach comes to mind:

from enum import Enum

class Station(Enum):
    wien_westbahnhof = 1
    st_poelten = 2
    linz = 3
    wels = 4

This way we can refer to a train station as Station.wien_westbahnhof in our code and get the name of an enumeration member from the name property:

>>> Station.wien_westbahnhof.name
'wien_westbahnhof'

This is handy but not usable for interaction with users. A user wants to see a station written as Wien Westbahnhof or St. Pölten and not the lowercase version that has an underline instead of a space in it.

What we could do now is to make a mapping from enumeration members to their "real world" representation:

STATIONS = {
    Station.wien_westbahnhof: 'Wien Westbahnhof',
    Station.st_poelten: 'St. Pölten',
    ...
}

This is an easy solution but has several drawbacks: We have to maintain the mapping outside of the class but use it in class methods to translate between the representations and the real enumeration members:

from enum import Enum

class Station(Enum):
    wien_westbahnhof = 1
    st_poelten = 2
    linz = 3
    wels = 4

    @classmethod
    def from_name(cls, name):
        for station, station_name in STATIONS.items():
            if station_name == name:
                return station
        raise ValueError('{} is not a valid station name'.format(name))

    def to_name(self):
        return STATIONS[self.value]

We also cannot use the default value of getting an instance from a name:

>>> Station['Wien Westbahnhof']
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/lukas/.pyenv/versions/3.6.1/lib/python3.6/enum.py", line 327, in __getitem__
    return cls._member_map_[name]
KeyError: 'Wien Westbahnhof'

Obviously, this fails because the Enum class has no idea about our mapping. We would have to use the from_name() method instead but this now leaves us with three possibilities to get an instance of a station.

Enums, however, also support aliases for names. I.e. it's possible for two values to have different names, so what we want is to have a string alias for each of our stations that represents the station in a human readable way:

class Station(Enum):
    wien_westbahnhof = 1
    'Wien Westbahnhof' = 1

This doesn't work because we can't have spaces in attribute names under normal circumstances.

However, there is a second way to create Enums: We can use the functional API that allows us to specify the member names using tuple lists or dicts. This way we can use names with spaces:

Station = Enum(
    value='Station',
    names=[
        ('wien_westbahnhof', 1),
        ('Wien Westbahnhof', 1),
        ('st_poelten', 2),
        ('St. Pölten', 2),
        ('linz', 3),
        ('Linz', 3),
        ('wels', 4),
        ('Wels', 4),
    ]
)

This way we can use the alias and the default name to get members. There is no no need for a translation table and separate methods anymore.

>>> Station['wien_westbahnhof']
<Station.wien_westbahnhof: 1>
>>> Station['Wien Westbahnhof']
<Station.wien_westbahnhof: 1>

We still have the problem that the default name and not the alias is returned as the value of the name property:

>>> Station['Wien Westbahnhof'].name
'wien_westbahnhof'

We can easily solve this by switching the order of alias and the default name in the list:

Station = Enum(
    value='Station',
    names=[
        ('Wien Westbahnhof', 1),
        ('wien_westbahnhof', 1),
        ('St. Pölten', 2),
        ('st_poelten', 2),
        ('Linz', 3),
        ('linz', 3),
        ('Wels', 4),
        ('wels', 4),
    ]
)

This way the human readable version is returned for the name property:

>>> Station['Wien Westbahnhof']
<Station.Wien Westbahnhof: 1>
>>> Station['wien_westbahnhof']
<Station.Wien Westbahnhof: 1>
>>> Station['wien_westbahnhof'].name
'Wien Westbahnhof'

The only drawback left is the duplicate values. By using a dict and some itertools magic we can also solve this problem:

import itertools

_STATIONS = {
    1: ['Wien Westbahnhof', 'WIEN_WESTBAHNHOF'],
    2: ['St. Pölten',       'ST_POELTEN'],
    3: ['Linz',             'LINZ'],
    4: ['Wels',             'WELS'],
}
Station = Enum(
    value='Station',
    names=itertools.chain.from_iterable(
        itertools.product(v, [k]) for k, v in _STATIONS.items()
    )
)

While this solution introduces a mapping again, it's also possible to use the dict directly without an intermediate variable.

Related Posts: