🔢 Enumerations in Python

python

Enumerations provide a developer with a restricted set of values.

Introduction

In many scenarios, you want a developer to pick one value from a well-defined set. Examples include the colors of a traffic light, HTTP methods, days of the week, or HTTP response status codes. To express this relationship in Python, you should use enumerations. Enumerations (enums) are a construct that let you define a list of valid values, and developers pick the specific value they want. Python first supported enumerations in Python 3.4 via the enum module.

Let's suppose I represent the traffic lights as a Python tuple:

# Note: use UPPER_CASE variable names to denote constant/immutable values
TRAFFIC_LIGHTS = ("Red", "Yellow", "Green")

What does this tuple communicate to other developers?

  • This collection is immutable.
  • They can iterate over this collection to get all the lights.
  • They can retrieve a specific light through static indexing.

The immutability and retrieval properties are important for my application. I don't want to add or subtract any lights at runtime. Retrieval lets me choose just one light, but it is a bit clunky.

TRAFFIC_LIGHTS[1]

This unfortunately does not communicate intent. Every time a developer sees this, they must remember that 1 means Yellow. Constantly correlating numbers to lights wastes time. This is fragile and will invariably cause mistakes. To combat this, I'll make aliases for each of these:

RED = "Red"
YELLOW = "Yellow"
GREEN = "Green"
TRAFFIC_LIGHTS = (RED, YELLOW, GREEN)

That's a bit more code, and still doesn't make it any easier to index into that tuple. Furthermore, there is still a lingering issue in calling code. Consider a function that performs an action based on the light:

def act(light: str):
    ...

Future developers would come across code like this:

act(TRAFFIC_LIGHTS[0])
act(RED)

Or:

act("Red")
# Definitely wrong
act("Red Yellow Green")

And here lies the crux of the problem. On the happy path, a developer can use the predefined variables. But if somebody accidentally were to use the wrong light, you soon get unwanted behavior. You want to find a way to communicate that you want a very specific, restricted set of values in specific locations. This is where enums shine.

Enum

Here's an example of Python's enumeration, Enum, in action:

from enum import Enum
# Class syntax
class TrafficLight(Enum):
    RED = "Red"
    YELLOW = "Yellow"
    GREEN = "Green"
# Functional syntax
TrafficLight = Enum('TrafficLight', ['RED', 'YELLOW', 'GREEN'])

To access specific instances, you can just do:

TrafficLight.RED
TrafficLight.GREEN

If you wanted to print out all the values of the enumeration, you can simply iterate over the enumeration:

for option_number, light in enumerate(TrafficLight, start=1):
    print(f"Option {option_number}: {light.value}")
# Option 1: Red
# Option 2: Yellow
# Option 3: Green

You can also get the name and value:

light = TrafficLight.RED
print(light.name)   # RED
print(light.value)  # Red

Finally, you can communicate your intent in functions that use this Enum:

def act(light: TrafficLight):
    ...

This tells all the developers looking at this function that they should be passing in a TrafficLight enumeration, and not just any old string. It becomes much harder to introduce typos or incorrect values. IDEs and type checkers can also help catch mistakes.

Enum Features

  • Iteration: Loop over all members.
  • Comparison: Enums are singletons, so TrafficLight.RED is TrafficLight.RED is True.
  • Uniqueness: Each member is unique.
  • Type safety: Functions can require a specific enum type.

When Not to Use

Enumerations are great for communicating a static set of choices for users. You don't want to use them where your options are determined at runtime, as you lose a lot of their benefits around communicating intent and tooling. If you find yourself in this situation, a dictionary which offers a natural mapping between two values that can be changed at runtime would be a better choice. You will need to perform membership checks if you need to restrict what values a user can select, though.

Advanced Usage

Automatic Values

For some enumerations, you might want to explicitly specify that you don't care about the value that the enumeration is tied to. This tells users that they should not rely on these values. For this, you can use the auto function.

from enum import auto, Enum
class TrafficLight(Enum):
    RED = auto()
    YELLOW = auto()
    GREEN = auto()
print(list(TrafficLight))
# [<TrafficLight.RED: 1>, <TrafficLight.YELLOW: 2>, <TrafficLight.GREEN: 3>]

By default, auto will select monotonically increasing values (1, 2, 3...). If you would like to control what values are set, you should implement a _generate_next_value_ function:

from enum import auto, Enum
class TrafficLight(Enum):
    def _generate_next_value_(name, start, count, last_values):
        return name.capitalize()
    RED = auto()
    YELLOW = auto()
    GREEN = auto()
print(list(TrafficLight))
# [<TrafficLight.RED: 'Red'>, <TrafficLight.YELLOW: 'Yellow'>, <TrafficLight.GREEN: 'Green'>]

Flags and Bitwise Operations

Sometimes you want to represent a combination of enum values, such as multiple permissions or statuses. The Flag and IntFlag classes allow you to combine enum members using bitwise operations.

from enum import auto, Flag
class TrafficLight(Flag):
    RED = auto()
    YELLOW = auto()
    GREEN = auto()

This lets you perform bitwise operations to combine traffic lights or check if certain lights are present.

stop_signs = TrafficLight.RED | TrafficLight.YELLOW
if stop_signs & TrafficLight.RED:
    print("You should wait!")
if stop_signs ^ TrafficLight.GREEN:
    print("You can go!")

Unique

One great feature of enumerations is the ability to alias values. However, there are cases where you want to force uniqueness on the values. Perhaps you are relying on the enumeration to always contain a set number of values, or perhaps it messes with some of the string representations that are shown to customers. No matter the case, if you want to preserve uniqueness in your Enum, simply add a @unique decorator.

from enum import Enum, unique
@unique
class TrafficLight(Enum):
    RED = "Red"
    YELLOW = "Yellow"
    GREEN = "Green"

If you accidentally assign the same value to two members, Python will raise a ValueError.

Integer Conversion

There are two more special case enumerations called IntEnum and IntFlag. These map to Enum and Flag, respectively, but allow degradation to raw integers for comparison. This can be useful for interoperability with C libraries or protocols, but be careful: comparing enums and integers directly can lead to subtle bugs if the enum values change.

from enum import IntEnum
class Status(IntEnum):
    OK = 200
    NOT_FOUND = 404
 
print(Status.OK == 200)  # True

Enum in Practice

Enumerations are widely used in Python's standard library and third-party packages. Examples include:

Enums also work well with type hints and static analysis tools, making your code more robust and self-documenting.

Summary

Enumerations are simple, and often overlooked as a powerful communication method. Any time that you want to represent a single value from a static collection of values, an enumeration should be your go-to user-defined type. It's easy to define and use them. They offer a wealth of operations, including iteration, bitwise operations, and control over uniqueness.

Further reading: