Month: February 2019

TL; DR – If you just want to see how to make the following work with string enums:

FooEnum.BAR == "Bar"
# True, without having to say `FooEnum.BAR.value`

… scroll down to the Trick™ section. Otherwise keep reading.

---

You might have already seen an application that used string literals for common constant values. For example, in an app with business objects that can be in different states, the code such as the following can be found:

if obj.state == "Active":
    # do something

...

if all_done():
  obj.state = "Completed"

Object states are represented by strings such as "Open", "Active", and "Completed", and there are many of these scattered around the code. Needless to say this implementation is not the best – it is susceptible to typos, and renaming a state requires a find & replace operation that can never go wrong (right?). A better approach is thus to store state names into constants (“constants” by convention at least), so that any future renamings can be done in a single place:

STATE_NEW = "Open"
STATE_ACTIVE = "Active"
STATE_DONE = "Completed"
...

if obj.state == STATUS_ACTIVE:
    # do something

...

if all_done():
  obj.state = STATE_DONE

If there are more than a few such constants defined in the application, it makes sense to group the related ones into namespaces. The most straightforward way is to defining them as class members:

class State:
    NEW = "Open"
    ACTIVE = "Active"
    DONE = "Completed"


if obj.state == State.ACTIVE:
    # etc.

Neat.

The State class has several drawback, however. Its members can be modified. Its members can be deleted. It is not iterable, thus compiling a list of all possible states is not elegant (one needs to peek into the class __dict__).

>>> State.NEW = "Completed"
>>> del State.ACTIVE
>>> list(
        (key, val) for key, val in State.__dict__.items()
        if not key.startswith('__')
    )
[('NEW', 'Open'), ('ACTIVE', 'Active'), ('DONE', 'Completed')]
>>> list(State)
...at's
TypeError: 'type' object is not iterable

Canonical solution: Enums

Starting with Python 3.4, the standard library provides the enum module that addresses these shortcomings (there is also a backport for older Python versions).

from enum import Enum

class State(Enum):
    NEW = "Open"
    ACTIVE = "Active"
    DONE = "Completed"

The only change is that the State class now inherits from Enum, suddenly making it more robust:

>>> State.NEW = "Completed"
# AttributeError: Cannot reassign members.
>>> del State.NEW
# AttributeError: State: cannot delete Enum member.
>>> list(State)
[<State.NEW: 'Open'>, <State.ACTIVE: 'Active'>, <State.DONE: 'Completed'>]

Each enum member is an object that has a name and a value:

>>> type(State.NEW)
<enum 'State'>
>>> State.NEW.name
'NEW'
>>> State.NEW.value
'Open'

There is a caveat, however – enum members can no longer be directly compared to string values:

>>> state = fetch_object_state(obj)  # assume "Open"
>>> State.NEW == state
False  # !!!

In order to work as expected, an enum member’s value must be compared:

>>> State.NEW.value == state
True

This is unfortunate, because the extra .value part makes the expression more verbose, and people might (rightfully) start complaining about readability. Not to mention that it represents a trap, it is too easy to forget about the .value suffix.

The standard library provides the IntEnum that makes the following work:

from enum import IntEnum

class Color(IntEnum):
    WHITE = 5
    BLACK = 10

>>> Color.WHITE == 5
True

Sadly, there is not “StringEnum” class, and it seems that you are on your own if you have string members. This sole reason can make some developers even think about ditching an enum altogether in favor of a plain class (first hand experience).

The trick™

And now for the primary motivation for this post. Thank you for reading it to here. 🙂

It is possible to use an enum while still preserving the convenience of a plain class when comparing the enum members to plain values. The trick is to subclass the type of enum members!

class State(str, Enum):  # <-- look, here
    NEW = "Open"
    ACTIVE = "Active"
    DONE = "Completed"

>>> State.NEW == "Open"
True
>>> State.NEW.value == "Open"
True

Even though this is described in the enum docs, one has to scroll down quite a lot towards the last quarter of the page to find it, thus you cannot blame yourself if you missed it the first time when you were just looking for a quick recipe.

With some creativity it is even possible to construct enums with types other than just the typical boring integers or strings. In a chess program, one could find the following enum useful to represent the corners of the board:

class Corner(tuple, Enum):
    TOP_LEFT = ('A', 8)
    TOP_RIGHT = ('H', 8)
    BOTTOM_LEFT = ('A', 1)
    BOTTOM_RIGHT = ('H', 1)

>>> rook_position = ('H', 8)
>>> is_top_corner = rook_position in (Corner.TOP_LEFT, Corner.TOP_RIGHT)
>>> is_top_corner
True

If you learned something new and found this trick useful, feel free to drop me a note. Thank you for reading!