Python Enum – How to create enums in Python

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

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

Let us begin!

Enum means “enumeration” and consists of a set of predefined named constants. These constants are usually related. Common examples include months of the year, days of the week, grades, order statuses, and tasks.

In short: an enum is essentially a collection of related constants, where each constant has a meaningful name associated with it.

In Python, you can create enumerations using the enumeration module (which we will do shortly!).

Why use enums

Using enums helps improve code clarity and maintainability. That is how:

• Enums improve code clarity by replacing magic numbers or strings with meaningful labels. They also make the code more self-documenting since the names of the enumeration members convey its purpose.
• Enums improve code maintainability by providing an easy way to define and manage related constants.
• By restricting variable assignments to only valid enum values, enums also ensure type safety.
• Enums make it easier to iterate and compare related constants.

Now let's create our first enum in Python.

We will create a TaskStatus enumeration that takes the following four names and values:

Image by author

First we import the Enum enum module class for creating enums.

Then we 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 this:

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

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

print(isinstance(TaskStatus.TODO,Enum))

Let's print all the members of the TaskStatus enum converting it to a list:

You should see the following result:

Output >>>

(<TaskStatus.TODO: 0>, <TaskStatus.IN_PROGRESS: 1>, <TaskStatus.DONE: 2>, <TaskStatus.ABANDONED: -1>)

All members of the enumeration have a name and a worth. Which means you can access enum members using their names, like TaskStatus.TODO. Or you access them for value, like TaskStatus(0).

Now that we have created a simple TaskStatus enum, let's learn how to perform simple tasks like iterating over enum members.

Iterating over enums

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

num_statuses = len(TaskStatus)
print(num_statuses)

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

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

Here is the result:

Output >>> 

TODO 0
IN_PROGRESS 1
DONE 2
ABANDONED -1

Sort on enumerations

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

• ALL: 0
• IN_PROGRESS: 1
• FACT: 2
• ABANDONED: -1

But you can also use the default order using the auto() auxiliary function. When you do this, if you have 'n' members in the enum, the values ​​assigned are 1 to n. But you can pass an initial value, say k, auto(k) so that the enumeration starts at k and goes up to k + n.

Let's modify the TaskStatus enumeration as shown:

from enum import Enum, auto

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

Now let's print 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 it TaskStatus enumeration 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

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

Image by author

we are going to 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 one update_status() Method that checks if the transition to the new state is valid given the current state. For invalid transitions, a ValueError exception is raised.

Here is an instance of the Task class: the “Write report” task with status ALL:

# 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 task state to IN_PROGRESS should work as it is 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 completed, 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 state to an invalid one, such as trying to update DONE to TODO, you will encounter the ValueError exception:

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

Here is the traceback of the ValueError raised due to an 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 create enums in Python by coding a simple TaskStatus enum. We learned how to access enum members and iterate over them.

Also, we learned how the default ordering works if you choose to use the auto() Helper function to set the values ​​of enumeration members. We then try to use the TaskStatus enum in a more useful 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. He enjoys working at the intersection of mathematics, programming, data science, and content creation. His areas of interest and expertise include DevOps, data science, and natural language processing. He likes to read, write, code and drink coffee! Currently, he is working to learn and share his knowledge with the developer community by creating tutorials, how-to guides, opinion pieces, and more.