Customize and Extend Python Enum Class

Summary: in this tutorial, you’ll learn how to customize and extend the custom Python enum classes.

Customize Python enum classes

Python enumerations are classes. It means that you can add methods to them, or implement the dunder methods to customize their behaviors.

The following example defines the PaymentStatus enumeration class:

from enum import Enum class PaymentStatus(Enum): PENDING = 1 COMPLETED = 2 REFUNDED = 3Code language: Python (python)

The PaymentStatus enumeration has three members: PENDING, COMPLETED, and REFUND.

The following displays the member of the PaymentStatus‘ member:

print(PaymentStatus.PENDING)Code language: Python (python)

It shows the following:

PaymentStatus.PENDINGCode language: Python (python)

To customize how the PaymentStatus member’s is represented in the string, you can implement the __str__ method. For example:

from enum import Enum class PaymentStatus(Enum): PENDING = 1 COMPLETED = 2 REFUNDED = 3 def __str__(self): return f'{self.name.lower()}({self.value})' print(PaymentStatus.PENDING)Code language: Python (python)

Now, it returns the following string:

pending(1)Code language: Python (python)

Implementing __eq__ method

The following attempts to compare a member of the PaymentStatus enum class with an integer:

if PaymentStatus.PENDING == 1: print('The payment is pending.')Code language: Python (python)

It shows nothing because the PaymentStatus.PENDING is not equal to integer 1.

To allow the comparison between PaymentStatus member and an integer, you can implement the __eq__ method like this:

from enum import Enum class PaymentStatus(Enum): PENDING = 1 COMPLETED = 2 REFUNDED = 3 def __str__(self): return f'{self.name.lower()}({self.value})' def __eq__(self, other): if isinstance(other, int): return self.value == other if isinstance(other, PaymentStatus): return self is other return False if PaymentStatus.PENDING == 1: print('The payment is pending.')Code language: Python (python)

In the __eq__ method:

• If the value to compare is an integer, it compares the value of the member with the integer.
• If the value to compare is an instance of the PaymentStatus enumeration, it compares the value with the member of the PaymentStatus member using the is operator.
• Otherwise, it returns False.

The program works as expected and returns the following output:

The payment is pending.Code language: Python (python)

Implementing __lt__ method

Suppose that you want the members of the PaymentStatus to follow have a sort order based on their value. And you also want to compare them with an integer.

To do that, you can implement the __lt__ method and use the @total_ordering decorator from the functools module:

from enum import Enum from functools import total_ordering @total_ordering class PaymentStatus(Enum): PENDING = 1 COMPLETED = 2 REFUNDED = 3 def __str__(self): return f'{self.name.lower()}({self.value})' def __eq__(self, other): if isinstance(other, int): return self.value == other if isinstance(other, PaymentStatus): return self is other return False def __lt__(self, other): if isinstance(other, int): return self.value < other if isinstance(other, PaymentStatus): return self.value < other.value return False # compare with an integer status = 1 if status < PaymentStatus.COMPLETED: print('The payment has not completed') # compare with another member status = PaymentStatus.PENDING if status >= PaymentStatus.COMPLETED: print('The payment is not pending')Code language: Python (python)

Implementing the __bool__ method

By default, all members of an enumeration are truthy. For example:

for member in PaymentStatus: print(member, bool(member))Code language: Python (python)

To customize this behavior, you need to implement the __bool__ method. Suppose you want the COMPLETED and REFUNDED members to be True while the PENDING to be False.

The following shows how to implement this logic:

from enum import Enum from functools import total_ordering @total_ordering class PaymentStatus(Enum): PENDING = 1 COMPLETED = 2 REFUNDED = 3 def __str__(self): return f'{self.name.lower()}({self.value})' def __eq__(self, other): if isinstance(other, int): return self.value == other if isinstance(other, PaymentStatus): return self is other return False def __lt__(self, other): if isinstance(other, int): return self.value < other if isinstance(other, PaymentStatus): return self.value < other.value return False def __bool__(self): if self is self.COMPLETED: return True return False for member in PaymentStatus: print(member, bool(member))Code language: Python (python)

The program output the following:

pending(1) False completed(2) True refunded(3) FalseCode language: Python (python)

Extend Python enum classes

Python doesn’t allow you to extend an enum class unless it has no member. However, this is not a limitation. Because you can define a base class that has methods but no member and then extend this base class. For example:

First, define the OrderedEnum base class that orders the members by their values:

from enum import Enum from functools import total_ordering @total_ordering class OrderedEnum(Enum): def __lt__(self, other): if isinstance(other, OrderedEnum): return self.value < other.value return NotImplementedCode language: Python (python)

Second, define the ApprovalStatus that extends the OrderedEnum class:

class ApprovalStatus(OrderedEnum): PENDING = 1 IN_PROGRESS = 2 APPROVED = 3Code language: Python (python)

Third, compare the members of the ApprovalStatus enum class:

status = ApprovalStatus(2) if status < ApprovalStatus.APPROVED: print('The request has not been approved.')Code language: Python (python)

Output:

The request has not been approved.Code language: Python (python)

Put it all together:

from enum import Enum from functools import total_ordering @total_ordering class OrderedEnum(Enum): def __lt__(self, other): if isinstance(other, OrderedEnum): return self.value < other.value return NotImplemented class ApprovalStatus(OrderedEnum): PENDING = 1 IN_PROGRESS = 2 APPROVED = 3 status = ApprovalStatus(2) if status < ApprovalStatus.APPROVED: print('The request has not been approved.')