Python Programming Tips #1

Using Enumerations

Python 3.4 introduced an enum module which at long last provides an enumeration type for Python. It has been worth the wait, since it has an excellent API. Furthermore, the module has been back-ported to every version of Python all the way to 2.4; so pretty well every Python programmer can take advantage of this useful module. (The back-ported module is called enum34.)

One ideal use is to replace a simple module constant (e.g., RED = 5), with an enum.Enum or an enum.IntEnum (this latter for when you actually care about the numeric value). However, out of the box, rather than writing, say, MyModule.RED you'd end up having to write, say, MyModule.State.RED. In this very short article I explain why this is and how to have enums that are drop-in replacements for simple module contants.

Let's take a look at a tiny example. First, a module that has some simple constants and a function, and second, a program that uses the module.

# Action.py
GO, STOP = range(2)

def perform(action):
    if action == GO:
        print("Go")
    elif action == STOP:
        print("Stop")

#!/usr/bin/env python3
import Action

Action.perform(Action.GO)
Action.perform(Action.STOP)

This works perfectly well and has been done for years. However, it isn't convenient for debugging since only the ints 0 and 1 would be shown.

The next example uses the same basic approach but is better for debugging.

# Action.py
GO, STOP = ("GO", "STOP")

def perform(action):
    if action == GO:
        print("Go")
    elif action == STOP:
        print("Stop")

#!/usr/bin/env python3
import Action

Action.perform(Action.GO)
Action.perform(Action.STOP)

I got the idea for this from Nick Coghlan.

Nonetheless, an enumeration would be nicer, and now that Python has an official enumeration module there's no excuse not to use it. Here's one way to do it (the enum module also has a function-based API).

# Action.py
import enum

class Action(enum.Enum):
    GO = 1
    STOP = 2

def perform(action):
    if action is Action.GO:
        print("Go")
    elif action is Action.STOP:
        print("Stop")

#!/usr/bin/env python3
import Action

Action.perform(Action.Action.GO)
Action.perform(Action.Action.STOP)

This works nicely, and puts the constants in their own namespace (the Action enum). Some people consider this the most Pythonic approach. However, this doesn't work as a drop-in replacement for simple constants, and some people find it a bit verbose.

Here's another version, this time using a technique that I discovered in Python 3.5's Lib/signals.py module.

# Action.py
import enum

class Action(enum.Enum):
    GO = 1
    STOP = 2
globals().update(Action.__members__)

def perform(action):
    if action is GO:
        print("Go")
    elif action is STOP:
        print("Stop")

#!/usr/bin/env python3
import Action

Action.perform(Action.GO)
Action.perform(Action.STOP)

The key difference here is the addition of one line immediately after the creation of the enumeration (and shown in bold). This line puts all the enumeration's members into the module's namespace. The result is that we can use GO and STOP directly in the module in which they are defined (compare this version's perform() function to the previous one), and in the same way as we use simple constants in modules that import them. This means that the program that's shown after the module with the enumeration has identical code to the first one shown, but is much nicer to debug because instead of getting plain ints we get enums which have both names and values.

Incidentally, the examples above assume that the actual enum values don't matter. If you want to use enums with numeric values as drop-in replacements for simple numeric constants, the enum.IntEnum class may be a better choice. Here's how.

# Action.py
import enum

class Action(enum.IntEnum):
    GO = 1
    STOP = 0
globals().update(Action.__members__)

def perform(action):
    if action is GO:
        print("Go")
    elif action is STOP:
        print("Stop")

#!/usr/bin/env python3
import Action

Action.perform(Action.GO)
Action.perform(Action.STOP)

This is almost identical the the previous example, but now the enumerations can also be used where ints are expected.

Although the globals().update(MyEnum.__members__) technique is used in the standard library, it is still a bit controversial: see the enum doc discussion. Personally, I have stopped using this technique, because it doesn't play well with Python linting tools e.g., flake8. My preferred style is now:

# Const.py
import enum

class ActionKind(enum.Enum):
    GO = 1
    STOP = 2

#!/usr/bin/env python3
import Action
from Const import *

Action.perform(ActionKind.GO)
Action.perform(ActionKind.STOP)

# Action.py
from Const import *

def perform(action):
    if action is ActionKind.GO:
        print("Go")
    elif action is ActionKind.STOP:
        print("Stop")

As you can see, I put all my constants and enumerations, into a single Const.py module and do from Const import *. My rule for this module is that it never contains anything mutable or that has side-effects. And for enum names I always use the convention NameKind so I never risk name collisions despite the import *, and always know that I have an enum rather than a simple NAME = value constant.

Incidentally, there is a third party advanced enum module, aenum, which provides some extra facilities compared with the standard library's enum module discussed here.

For more see Python Programming Tips

Top