Article Thumbnail

Why You Should Use More Enums In Python

A gentle introduction to enumerations in Python

Florian Dahlitz
10 min
June 8, 2020

Introduction

In this article, you will learn what enums are and how to create and work with them. Furthermore, you will learn why you should use them more often in your day to day coding.

What Is an Enum?

enum stands for enumeration and refers to a set of symbolic names, which are called enumeration members. These enum members are bound to unique, constant values. You can iterate over an enumeration and compare its members by identity (Python's is operator).

The following code snippet shows you a simple example of an enum Colour:

# colour.py

from enum import Enum


class Colour(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3

We imported the Enum class from Python's enum module. It serves as a base class for defining new enumerations in Python. Subsequently, a new enum called Colour is implemented having three enum members: RED, GREEN, and BLUE.

Note: Although the class syntax is used to define new enumerations, they aren't normal Python classes. If you want to know more about it, check out the How are Enums different? section in the module's documentation [1]:

Let's see how enums behave when used.

# previous colour.py code

c = Colour.RED
print(c)
print(c.name)
print(c.value)
print(c is Colour.RED)
print(c is Colour.BLUE)

We extended the colour.py script by creating a new instance of Colour.RED and assigning it to the variable c. Furthermore, we print the string representation of Colour.RED, its name and value. Additionally, we compare c's identity with Colour.RED and Colour.BLUE.

$ python colour.py
Colour.RED
RED
1
True
False

Running the script reveals that c is indeed an instance of Colour.RED with RED as its name and 1 as value.

Note: We used the is operator to compare the variable c with the different enum members. Keep in mind that enums can only be compared to enums and not to integers, even though the enum member values are integers [2].

Iterating Over the Members of an Enum

Enumerations have a special attribute called __members__, which is a read-only ordered mapping of names and members. Utilising __members__ allows you to iterate over an enum and print the members as well as their corresponding names.

# iterate.py

from enum import Enum


class Colour(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3


for name, member in Colour.__members__.items():
    print(name, member)
$ python iterate.py
RED Colour.RED
GREEN Colour.GREEN
BLUE Colour.BLUE

You might ask yourself why we did not something like:

for member in Colour:
    print(member.name, member)

For the example at hand, both approaches produce the same result. However, if you have an enumeration that has aliases, too, only the approach using __members__ will print the aliases as well. Checkout the following example:

# iterate_alias.py

from enum import Enum


class Colour(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3
    ALIAS_RED = 1


for name, member in Colour.__members__.items():
    print(name, member)

print("="*20)

for member in Colour:
    print(member.name, member)
$ python iterate_alias.py
RED Colour.RED
GREEN Colour.GREEN
BLUE Colour.BLUE
ALIAS_RED Colour.RED
====================
RED Colour.RED
GREEN Colour.GREEN
BLUE Colour.BLUE

Automatic Values

In the previous example, we assigned integers to the symbolic names RED, GREEN, and BLUE. If the exact values are not important, you can use the enum.auto() function. The function calls _generate_next_value_() internally and generates the values for you.

# auto.py

from enum import auto
from enum import Enum


class Colour(Enum):
    RED = auto()
    GREEN = auto()
    BLUE = auto()


c = Colour.RED
print(c.value)

It chooses a suited value for each enum member, which will (most of the time) be the integers we used before.

$ python auto.py
1

However, the _generate_next_value_() function can be overwritten to generate new values the way you like:

# overwritten_next_values.py

from enum import auto
from enum import Enum


class AutoName(Enum):
    def _generate_next_value_(name, start, count, last_values):
        if len(last_values) > 0:
            return last_values[-1] * 2
        return 2


class Colour(AutoName):
    RED = auto()
    GREEN = auto()
    BLUE = auto()


c = Colour.RED
g = Colour.GREEN
b = Colour.BLUE
print(c.value)
print(g.value)
print(b.value)
$ python overwritte_next_values.py
2
4
8

Extending an Enum

Being Python classes, enums can have any (special) methods just like all other classes. Consider the following example.

# extending.py

from enum import Enum


class Colour(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3

    def __str__(self):
        return self.name

    def colorize(self):
        return f"Let's paint everything in {self.name.lower()}"


c = Colour.RED
print(c)
print(c.colorize())

We extended the Colour enum by a new method colorize() printing a string generated based on the member's value. Furthermore, we overwrite the __str__() dunder method to return the member's name if called.

$ python extending.py
RED
Let's paint everything in red

Kinds of Enums in Python

Besides Enum, Python provides three derived enumerations out of the box:

  • IntEnum
  • IntFlag
  • Flag

We will have a look at all three of them. Keep in mind that you are free to implement your own derived enumerations based on Enum. Implementing your own enumeration will not be covered in the article.

IntEnum

We already know that we can compare enum members using Python's identity operator. However, the Enum class does not provide ordered comparisons even though integers are used as values for the enumeration members. Let's have a look at the following example.

# comparison.py

from enum import Enum


class Colour(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3


r = Colour.RED
b = Colour.GREEN
print(r < b)

Executing the script at hand results in a TypeError.

$ python comparison.py
Traceback (most recent call last):
  File "/home/florian/workspace/python/why-you-should-use-more-enums-in-python-article-snippets/comparison.py", line 14, in <module>
    print(r < b)
TypeError: '<' not supported between instances of 'Colour' and 'Colour'

The only thing you can do is making use of equality comparisons like == and !=. Additionally, comparing enum members with any non-enumeration value is not supported.

However, the derived enumeration IntEnum does provide ordered comparisons as it is also a subclass of int. In order to make our example work, we need to import the IntEnum class instead of Enum and derive Colour from it. We do not need to change anything else.

# comparison.py

from enum import IntEnum


class Colour(IntEnum):
    ...
$ python comparison.py
True

IntFlag

The IntFlag class is pretty similar to the IntEnum class with the exception that is also supports bitwise operations. With supporting bitwise operations I mean that it is possible to combine two enum members resulting in an IntFlag member, too. All other operations on an IntFlag member will result in the loss of the IntFlag membership.

Let's have a look at an example. Assume that we grant permissions to users so that they can read, write and/or execute a certain file. We create an enumeration Permission with the members R (read permission), W (write permission), and X (execute permission) respectively.

If we have a user, who should have read and write permissions for a certain file, we can combine both using the | operator.

# permissions.py

from enum import IntFlag


class Permission(IntFlag):
    R = 4
    W = 2
    X = 1


RW = Permission.R | Permission.W
print(RW)
print(Permission.R + Permission.W)
print(Permission.R in RW)
$ python permissions.py
Permissions.R|W
6
True

Flag

The Flag class does also provide support for bitwise operations but does not inherit from int. In fact, it is like Enum but with support for the bitwise operations.

If we take the Colour enum from the beginning, we could easily mix the colour white based on the other three colours.

# colour_flag.py

from enum import auto
from enum import Flag


class Colour(Flag):
    RED = auto()
    GREEN = auto()
    BLUE = auto()
    WHITE = RED | GREEN | BLUE


print(Colour.WHITE.name, Colour.WHITE.value)
$ python colour_flag.py
WHITE 7

Why Do I Need Enums?

At this point, we have an understanding of what enums are and how we can create them in Python. Furthermore, we are able to compare and work with them. However, we still do not know why we need and should use enumerations more often.

The examples we had a look at were pretty simple. Although the Permission enumeration seems pretty useful, the Colour enum does not. Why would you use these enumerations in your code?

Have a look at the following Python snippet.

# response_code_magic_numbers.py

from http.client import HTTPResponse


def evaluate_response(response: HTTPResponse) -> str:
    if response.code() == 404:
        return "Not Found"
    elif response.code() == 502:
        return "???"
    elif response.code() == 400:
        return "???"
    else:
        return "Unknown Status Code"

We defined a function, which takes an HTTPResponse object and returns a string based on the status code of the supplied HTTPResponse object. You may know that 404 is the status code for Not Found, but do you know the meaning of 502 and 400?

These are only two less known status codes and much more are out there. It is hard to read and understand the code without a web search. This is where enumerations come into play. We can implement our own custom enumeration to lend more meaning to the code.

# http_code_enum.py

from enum import IntEnum


class HTTPCode(IntEnum):
    BAD_REQUEST = 400
    NOT_FOUND = 404
    BAD_GATEWAY = 502

Here an IntEnum is used as we want to be able to compare members of it with integers. Now, the function from before looks like this:

# response_code.py

from http_code_enum import HTTPCode
from http.client import HTTPResponse


def evaluate_response(response: HTTPResponse) -> str:
    if response.code() == HTTPCode.NOT_FOUND:
        return "Not Found"
    elif response.code() == HTTPCode.BAD_GATEWAY:
        return "???"
    elif response.code() == HTTPCode.BAD_REQUEST:
        return "???"
    else:
        return "Unknown Status Code"

In essence, if you have magic numbers in your code, you should definitely consider to either assign them to a variable or group them together to an enumeration. This way your code's readability increases a lot. It is especially true if you want to write tests for your code.

Summary

Congratulations, you have made it through the article! While reading the article you learned what enums are and how you can create them in Python. Furthermore, you learned how to compare and work with them. You had a look at a few examples and understood why it is good practice to use enumerations once in a while.

I hope you enjoyed reading the article. Feel free to share it with your friends and colleagues! Do you have feedback? I am eager to hear it! You can contact me via the contact form or other resources listed in the contact section.

If you have not already, consider following me on Twitter, where I am @DahlitzF, or subscribing to my newsletter! Stay curious and keep coding!

References