Python Enum: How To Build Enumerations in Python

Image by Author

If you’ve programmed in a language like C++ or Java, you’ve likely used enums to create named constants. Enums are helpful when you have a variable that takes one of a fixed number of values—often related such as the days of the week, student grades, order status, and the like.

Python, however, does not have an explicit enum data type. But you can use the enum module in the Python standard library to create enumerations. And this tutorial will teach you how. 

Let’s get started!

Enum stands for “enumeration” and consists of a set of predefined named constants. These constants are often related. Common examples include months in a year, days of the week, grades, order and task statuses. 

To sum up: An enum is essentially a collection of related constants, where each constant has a meaningful name associated with it.

In Python, you can create enums using the enum module (which we’ll do shortly!).

Why Use Enums

Using enums helps improve code clarity and maintainability. Here’s’ how: 

• Enums enhance code clarity by replacing magic numbers or strings with meaningful labels. They also make the code more self-documenting as the names of enum members convey their purpose. 
• Enums improve code maintainability by providing a simple way to define and manage related constants. 
• By restricting variable assignments to only valid enum values, enums also ensure type safety. 
• Enums facilitate easy iteration over and comparison of related constants.

Now let’s create our first enumeration in Python.

We’ll create a TaskStatus enum that takes the following four names and values:

Image by Author

First, we import the Enum class from the enum module to create enums. 

We then define a new class TaskStatus that inherits from Enum to create an enumeration. Each member is defined with a unique name and an optional value like so:

from enum import Enum
	 
class TaskStatus(Enum):
    TODO = 0
    IN_PROGRESS = 1
    DONE = 2
    ABANDONED = -1

All enum members we create are instances of the Enum class. You may verify it by calling the isinstance() function as shown:

print(isinstance(TaskStatus.TODO,Enum))

Let’s print out all the members in the TaskStatus enum by casting it into a list:

You should see the following output:

Output >>>

[<TaskStatus.TODO: 0>, <TaskStatus.IN_PROGRESS: 1>, <TaskStatus.DONE: 2>, <TaskStatus.ABANDONED: -1>]

All enum members have a name and a value. Meaning you can access enum members using their names, like TaskStatus.TODO. Or you access them by value, like TaskStatus(0).

Now that we’ve created a simple TaskStatus enum, let’s learn how to perform simple tasks such as iterating over the enum members.

Iterating Over Enums

In Python, you can work with enums pretty much the same way you work with any iterable. For example, you can use the len() function to count the number of enum members:

num_statuses = len(TaskStatus)
print(num_statuses)

You can also iterate over enums just the way you would over a Python iterable such as list. In the following for loop, we access both the name and value of each enum member and print them out:

for status in TaskStatus:
    print(status.name, status.value)

Here’s the output:

Output >>> 

TODO 0
IN_PROGRESS 1
DONE 2
ABANDONED -1

Ordering in Enums

In the example, the status and the corresponding numeric value are as follows:

• TODO: 0
• IN_PROGRESS: 1
• DONE: 2
• ABANDONED: -1

But you can also use the default ordering by using the auto() helper function. When you do so, if you have ‘n’ members in the enum, the values assigned are 1 through n. But you can pass in a start value, say k, auto(k) for the enumeration to start at k and go up to k + n.

Let’s modify the TaskStatus enum as shown:

from enum import Enum, auto

class TaskStatus(Enum):
    TODO = auto()
    IN_PROGRESS = auto()
    DONE = auto()
    ABANDONED = auto()

Now let’s print out the members:

We see that the values are 1 to 4 as expected:

Output >>>

[<TaskStatus.TODO: 1>, <TaskStatus.IN_PROGRESS: 2>, <TaskStatus.DONE: 3>, <TaskStatus.ABANDONED: 4>]

Now let’s build on the TaskStatus enum that we have. Create a task.py file in your working directory with the following version of the enum:

# task.py

from enum import Enum

class TaskState(Enum):
    TODO = 0
    IN_PROGRESS = 1
    DONE = 2
    ABANDONED = -1

Say we have a task with a name and a current status. And the valid transitions between states are as shown:

Image by Author

Let’s create a Task class:

class Task:
    def __init__(self, name, state):
        self.name = name
        self.state = state
    
    def update_state(self, new_state):
        # Define valid state transitions based on the current state
        valid_transitions = {
        	TaskState.TODO: [TaskState.IN_PROGRESS, TaskState.ABANDONED],
        	TaskState.IN_PROGRESS: [TaskState.DONE, TaskState.ABANDONED],
        	TaskState.DONE: [],
        	TaskState.ABANDONED: []
    	}
   	 
        # Check if the new state is a valid transition from the current state
        if new_state in valid_transitions[self.state]:
            self.state = new_state
        else:
            raise ValueError(f"Invalid state transition from {self.state.name} to {new_state.name}")

We have an update_status() method that checks if the transition to the new state is valid given the current status. For invalid transitions, a ValueError exception  is raised.

Here’s an Instance of the Task class: the “Write Report” task with status TODO:

# Create a new task with the initial state "To Do"
task = Task("Write Report", TaskState.TODO)

# Print the task details
print(f"Task Name: {task.name}")
print(f"Current State: {task.state.name}")

Output >>>
Task Name: Write Report
Current State: TODO

Updating the status of the task to IN_PROGRESS should work as it’s a valid state transition:

# Update the task state to "In Progress"
task.update_state(TaskState.IN_PROGRESS)
print(f"Updated State: {task.state.name}")

Output >>> Updated State: IN_PROGRESS

And once the task is complete, we can update its status to DONE:

# Update the task state to "DONE"
task.update_state(TaskState.DONE)
print(f"Updated State: {task.state.name}")

Output >>> Updated State: DONE

But if you try to update the status to an invalid one, such as trying to update DONE to TODO, you’ll run into ValueError exception:

# Attempt to update the task state to an invalid state
task.update_state(TaskState.TODO)

Here’s the traceback of the ValueError raised because of invalid state transition from DONE to TODO:

Traceback (most recent call last):
  File "/home/balapriya/enums/task.py", line 46, in 
	task.update_state(TaskState.TODO)
  File "/home/balapriya/enums/task.py", line 30, in update_state
	raise ValueError(f"Invalid state transition from {self.state.name} to {new_state.name}")
ValueError: Invalid state transition from DONE to TODO

In this tutorial, we learned how to build enumerations in Python by coding a simple TaskStatus enum. We learned how to access enum members and iterate over them. 

Also, we learned how default ordering works if you choose to use the auto() helper function to set the values for enum members. We then tried using the TaskStatus enum in a more helpful example. 

You can find the code examples on GitHub. I’ll see you all soon in another Python tutorial. Until then, happy coding!
 
 

Bala Priya C is a developer and technical writer from India. She likes working at the intersection of math, programming, data science, and content creation. Her areas of interest and expertise include DevOps, data science, and natural language processing. She enjoys reading, writing, coding, and coffee! Currently, she’s working on learning and sharing her knowledge with the developer community by authoring tutorials, how-to guides, opinion pieces, and more.