Supercharge Your Enums: Cleaner Code with Hidden Features

I see a lot of articles suggesting you use enums by mostly restating the Python documentation. Unfortunately, I feel this leaves readers without crutial practical advice, which I’d like to pass on here.

This is especially true since most of the projects I’ve worked on, and developers I’ve coded with, don’t seem to know this, leading to more complex code, including values and behavior that are tightly coupled in the business concepts scattered about in separate code files.

First, let’s review enum fundamentals:

Enums in a nutshell

If you’re not aware of enums, they were added to Python in version 3.4 and represent a way to communicate, among other things, a reduced set of options.

For example, you can communicate the status of tasks:

from enum import Enum

class TaskStatus(Enum):  
    PENDING = 'pending'  
    IN_PROGRESS = 'in_progress'  
    COMPLETED = 'completed'  
    CANCELLED = 'cancelled'

The four lines of code beneath the TaskStatus class definition define the enum “members,” or the specific options for this class.

This code communicates the fact that there are only four statuses a task can have.

This is very useful information, as it clearly shows the options available, but it is as deep as many developers go with enums. They don’t realize how much more enums can do besides holding constant values.

For example, many don’t know how easy it can be to select an enum member.

Enums can select themselves

I see a lot of code that complicates selecting enum members, like this:

def change_task_status(task_id: str, status: str):  
    task = database.get_task_by_id(task_id)  
    for member in TaskStatus:  
        if member.value == status:  
            task.status = member  
            database.update_task(task)

Instead, enum classes are smart enough to select members from their values (the things on the right side of the equal sign)You can also select members by their names (the left side of the equal sign) with square brackets, TaskStatus['PENDING'] . :

>>> TaskStatus('pending')
<TaskStatus.PENDING: 'pending'>

This means that you could simplify the code above like thisBe aware that if the status string does not match one of the enum member’s values, it’ll raise a ValueError . :

def change_task_status(task_id: str, status: str):  
    task = database.get_task_by_id(task_id)  
    task.status = TaskStatus(status)  
    database.update_task(task)

But enums are not just static values. They can have behavior and data associated with them too.

Enum members are objects too

The thing about enums that many people are missing is that they are objects too.

For example, I recently worked on a project that would have had this after the TaskStatus class to connect a description to each enum member:

STATUS_TO_DESCRIPTION_MAP = {  
    TaskStatus.PENDING: "Task is pending",  
    TaskStatus.IN_PROGRESS: "Task is in progress",  
    TaskStatus.COMPLETED: "Task is completed",  
    TaskStatus.CANCELLED: "Task is cancelled"  
}

But here’s the thing, you can add it in the enum!

Granted, it takes a little bit of work, but here’s how I would do itIf we weren’t using the enum’s value to select a member, we could make this simpler by editing the __init__ method instead, like they do in the docs. :

class TaskStatus(Enum):  
    PENDING = "pending", 'Task is pending'  
    IN_PROGRESS = "in_progress", 'Task is in progress'  
    COMPLETED = "completed", 'Task is completed'  
    CANCELLED = "cancelled", 'Task is cancelled' 
 
    def __new__(cls, value, description):  
        obj = object.__new__(cls)  
        obj._value_ = value  
        obj.description = description  
        return obj

This means that whenever a TaskStatus member is created, it keeps its original value but also adds a new attribute, description.

This means that a TaskStatus member would behave like this:

>>> completed = TaskStatus.COMPLETED
>>> completed.value
'completed'
>>> completed.description
'Task is completed'
>>> completed.name
'COMPLETED'

On top of that, you can define methods that interact with the enum members.

Let’s add what the business expects would be the next status for each member:

class TaskStatus(Enum):  
    PENDING = "pending", "Task is pending"  
    IN_PROGRESS = "in_progress", "Task is in progress"  
    COMPLETED = "completed", "Task is completed"  
    CANCELLED = "cancelled", "Task is cancelled"  
  
    def __new__(cls, value, description):  
        ...
  
    @property  
    def expected_next_status(self):  
        if self == TaskStatus.PENDING:  
            return TaskStatus.IN_PROGRESS  
        elif self == TaskStatus.IN_PROGRESS:  
            return TaskStatus.COMPLETED  
        else:  # Task is completed or cancelled  
            return self

Now, each TaskStatus member “knows” what status is expected to be next:

>>> TaskStatus.PENDING.expected_next_status
<TaskStatus.IN_PROGRESS: 'in_progress'>
>>> TaskStatus.CANCELLED.expected_next_status
<TaskStatus.CANCELLED: 'cancelled'>

You could use this in a task detail view:

def task_details(task_id: str):  
    task = database.get_task_by_id(task_id)  
    return {  
        "id": task.id,  
        "title": task.title,  
        "status": task.status.value,  
        "expected_next_status": task.status.expected_next_status.value,  
    }

>>> task_details("task_id_123")
{
  'id': 'task_id_123', 
  'title': 'Sample Task', 
  'status': 'pending', 
  'expected_next_status': 'in_progress'
}

In conclusion

Python enums are more powerful than most developers realize, and I hope you might remember these great options.