Enumerations in Python

Python enumerations

Python enumerations or Enums can be used to couple related values together in a type-safe way.

For example, here is an enumeration for directions:

from enum import Enum

class Direction(Enum):
    UP = 1
    DOWN = 2
    LEFT = 3
    RIGHT = 4

This makes it possible to call Direction.UP or Direction.DOWN in your program.

Notice that Python enumerations have nothing to do with the enumerate() function. If you came for enumerate() check out this article.

Python Enumerations in Detail

Python supports enumerations.

An enumeration is a group of symbolic names tied to constant values called enumeration members. So enumeration is just a type for a group of constant values.

To support enumerations in your Python project, import the Enum class from the enum module.

To create an enumeration, use the class syntax. Notice that enumerations are not normal Python classes, even though they are defined in a similar way.

As enumerations are meant to represent constant values, it is advisable to uppercase the enum member names.

For instance, let’s create an enumeration for the days of the week:

from enum import Enum

class Day(Enum):
    SUN = 1
    MON = 2
    TUE = 3
    WED = 4
    THU = 5
    FRI = 6
    SAT = 7

Let’s next inspect how the enumeration behaves.

Print Enum Member

You can either print the enumeration member directly, or you can access its name or value properties.

Here is how it works (using the Day enum we created earlier):

print(Day.MON)
print(Day.MON.name)
print(Day.MON.value)

Output:

Day.MON
MON
2
  • Printing the enumeration member as-is results in it being printed into the console as-is.
  • Printing the name of the enumeration member only prints the name of the constant.
  • Printing the value of the member prints the numeric value assigned to the constant.

These are the three basic ways you can access the enumeration object’s properties.

Comparing Enumeration Members

You can compare enumeration members by their values.

For example, let’s compare the members of the Day enum we created before:

print(Day.MON == Day.MON)
print(Day.TUE == Day.MON)
print(Day.WED != Day.MON)

Output:

True
False
True

Notice, however, that comparing order is not possible using Enums. In other words, operators <, >, <=, >= do not work.

Loop Through the Values of an Enum

You can loop through an enumeration using a for loop. And you do not need to create an enumeration object to do that. Instead, you can call the name of the enumeration directly.

For example, let’s print the names of the days of the Day enumeration:

for day in Day:
    print(day.name)

Output:

SUN
MON
TUE
WED
THU
FRI
SAT

Looping the enumeration only works if your enumeration does not contain aliases. In other words, if each enumeration member has a unique value. If there are multiple members with the same value, you need to use a different approach.

An enumeration has a built-in __members__ attribute. This attribute maps each enumeration member’s name to the member. In other words, all the members can be found in this dictionary, regardless of the member value.

Let’s create an Enum that carries different members with equal values:

from enum import Enum

class Day(Enum):
    SUN = 1
    MON = 2
    TUE = 3

    FIRST_DAY = 1
    SECOND_DAY = 2
    THIRD_DAY = 3

To loop through this Enum, use the __members__ dictionary, and loop through it like a regular dictionary.

Here is how:

for name, member in Day.__members__.items():
    print(f"{name}: {member}")

Output:

SUN: Day.SUN
MON: Day.MON
TUE: Day.TUE

FIRST_DAY: Day.SUN
SECOND_DAY: Day.MON
THIRD_DAY: Day.TUE

If you used the regular loop here, you would only see the first three enumeration members:

for day in Day:
    print(day.name)

Output:

SUN
MON
TUE

Use Enum Members in Dictionaries

Python enumerations are hashable. This means you can use the enumeration members in dictionaries.

Let’s bring back the Day enumeration you saw earlier.

from enum import Enum

class Day(Enum):
    SUN = 1
    MON = 2
    TUE = 3
    WED = 4
    THU = 5
    FRI = 6
    SAT = 7

Here is an example of mapping the Day members to a textual descriptions into a dictionary.

schedule = {
    Day.SUN: "Relax",
    Day.MON: "Work",
    Day.TUE: "Meetings",
    Day.WED: "Mid-week party",
    Day.THU: "Morning Workout",
    Day.FRI: "Fine-dining",
    Day.SAT: "Road Trip"
}

print(schedule[Day.FRI])

Output:

Fine-dining

Automatic Enum Values

In the previous examples, you have seen weekday enumeration. In this enumeration it makes sense to give the numbers 1...7 as the values of the days.

If you are creating an enumeration of something unrelated to numbers, you don’t have to assign the number values 1...n to it.

In this case, you can let Python assign the values on your behalf. To do this, use the enum.auto() function. Remember to import the auto() function from the enum module.

For instance, let’s create a Fruit enumeration. We don’t care about the value of the enumeration member, because numbering a fruit is meaningless. Thus, we can use the auto() function to generate values automatically:

from enum import Enum, auto

class Fruit(Enum):
    APPLE = auto()
    PINEAPPLE = auto()
    ORANGE = auto()
    PEAR = auto()

But it turns out that most of the time, the auto() function assigns uses similar integer values as you would.

Let’s print the enumeration member values of the fruits:

for fruit in Fruit:
    print(fruit.value)

Output:

1
2
3
4

Custom Automatic Enumeration Values

In the previous chapter, you saw how you can auto-generate values for the enumeration members.

But you can also override the method that generates the automatic values. This method is called generate_next_value.

Without digging into too much details, let’s generate negative values for the enum members by overriding thegenerate_next_value method:

from enum import Enum, auto

class Direction(Enum):
    def _generate_next_value_(name, start, count, last_values):
        if len(last_values) > 0:
            return last_values[-1] - 1
        return -1
    
    UP = auto()
    RIGHT = auto()
    DOWN = auto()
    LEFT = auto()

print(Direction.UP.value)
print(Direction.RIGHT.value)
print(Direction.DOWN.value)
print(Direction.LEFT.value)

Output:

-1
-2
-3
-4

Extending Enums

Python enumerations have things in common with regular classes. One of them is the ability to extend the enumerations with methods.

For example, let’s make it possible for the enumeration members to describe themselves:

from enum import Enum

class Day(Enum):
    SUN = 1
    MON = 2
    TUE = 3
    WED = 4
    THU = 5
    FRI = 6
    SAT = 7
    
    def describe(self):
        return f"Today is {self.name}."

print(Day.MON.describe())

Output:

Today is MON.

Enumeration Types in Python

In addition to the basic Enum, there are three subtypes of it.

These are:

  • IntEnum
  • IntFlag
  • Flag

Let’s take a quick look at each.

IntEnum

Comparing enumeration member values is possible using the identity operator (== and !=).

But you cannot compare the order of enumeration members. This can be counterintuitive because the value is an integer:

print(Day.MON < Day.TUE)

Output:

TypeError: '<' not supported between instances of 'Day' and 'Day'

But there is a solution. Use a derivative of the Enum class called IntEnum:

from enum import IntEnum

class Day(IntEnum):
    SUN = 1
    MON = 2
    TUE = 3
    WED = 4
    THU = 5
    FRI = 6
    SAT = 7

print(Day.MON < Day.TUE) # Prints True

IntFlag

The IntFlag enumeration class is similar to the IntEnum class. The main difference is it only supports bitwise operations (&, |, ^, ~). In other words, you can combine IntFlag enumeration members to create new IntFlag members.

For example, let’s create a Permission (Read, Write, Execute) IntFlag. This can represent possible permissions granted for new users in a workspace for example.

from enum import IntFlag

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

Let’s create a new IntFlag object for writing and reading, but not executing.

RW = Permission.R | Permission.W

Now RW is a new Permission object. Let’s play with it to understand how it works:

print(RW)                   # Prints Permission.R|W
print(RW.value)             # Prints 6 (4|2 = 0100|0010 = 0110 = 6)
print(Permission.R in RW)   # Prints True

With an IntFlag object, it is also possible to create new ones by combining with integers.

For example, let’s create an RWX object by combining Permission.R (4) with 2 and 1:

print(Permission.R | 2 | 1) # Permission.R|W|X

Here, the IntFlag object automatically converts the 2 and 1 to Permissions.

Flag

The Flag enumeration class is also meant for bitwise operations.

Let’s create a Flag that represents RGB colors:

from enum import Flag, auto

class Color(Flag):
    R = auto()
    G = auto()
    B = auto()

Now let’s create white color by combining R,G, and B:

WHITE = Color.R | Color.G | Color.B

print(WHITE)

Output:

Color.B|G|R

The Flag behaves similar to IntFlag. But unlike IntFlag, the Flag cannot be combined nor compared with other Flag enumeration objects or Int.

For example, this combination that would be possible using IntFlag, fails with Flag:

Color.R | 4
# TypeError: unsupported operand type(s) for |: 'Color' and 'int'

So there you have it. An introduction to Python enumerations and its derivatives. Let’s finally take a look at when you benefit from using an Enum in Python.

An Example Use Case for Python Enumeration

Say you have a program that controls the direction of a game character. This character can move up, down, left, and right.

Here is a function that handles actions based on a given direction:

def handleDirection(direction):
    if direction == "UP":
        print("Moving up")
    elif direction == "DOWN":
        print("Moving down")
    elif direction == "LEFT":
        print("Moving left")
    elif direction == "RIGHT":
        print("Moving right")
    else:
        print("Not moving")

This function expects one of the following strings UP, DOWN, LEFT, RIGHT as its arguments.

In case the input isn’t one of those, Not moving is printed.

This function is valid in Python and works as expected. The problem is that it is prone to bugs.

Let’s see what happens if we make a small mistake when calling the function:

handleDirection("up")

Output:

Not moving

This happens because we spelled “up” incorrectly. It should have been "UP", not "up". Here is the proble. There is no way to tell what direction names you can use other than going back to the definition of handleDirection.

This is bad practice.

Here is where Python enumerations help.

First, let’s declare an enumeration for directions:

from enum import Enum

class Direction(Enum):
    UP = 1
    DOWN = 2
    LEFT = 3
    RIGHT = 4

Instead of using mystic strings, you now have Direction—a custom type that specifies the four possible directions.

Now, you can pass Direction enumeration object as an argument to the handleDirection function instead of using strings.

def handleDirection(direction):
    if direction == Direction.UP:
        print("Moving up")
    elif direction == Direction.DOWN:
        print("Moving down")
    elif direction == Direction.LEFT:
        print("Moving left")
    elif direction == Direction.RIGHT:
        print("Moving right")
    else:
        print("Not moving")

You can now call handleDirection this way:

handleDirection(Direction.UP)

Output:

Moving up

This is handy. You don’t need to go back to the handleDirection function to check what possible directions you can use. Instead, you can either check the Direction enumeration or see what the code editor suggests:

Enumerations in Python

This is way cleaner approahc than using magic strings.

Conclusion

Today you learned about Python enumerations.

You now know how to create one, how to loop through it, and how to make comparisons with it.

You also learned how you can customize enums with your own methods, and how to use some common derivatives of enums.

To take home, Python enumerations are special types of classes that allow you group constant values together. These can be useful from time to time.

Thanks for reading. I hope you enjoy it.

Happy coding!

Further Reading

50 Python Interview Questions

Enumerate Function in Python

Share on facebook
Facebook
Share on google
Google+
Share on twitter
Twitter
Share on linkedin
LinkedIn
Share on pinterest
Pinterest

Leave a Comment

Your email address will not be published.