A Guide to Using Enum in Python for Finite Data States

This post is going to explore the Python enum module, a powerful tool for handling sets of data that don’t change.

Not only is this integral for generalized software development, but is something you’ll also find commonly featured in games – making it an important skill and foundation to learn regardless of your personal coding goals!

Let’s dive in, and explore this fascinating topic of using enum in Python!

Prerequisites: The reader is expected to have a working knowledge of the Python programming language. Familiarity with how to write classes will also be helpful since Python enums are implemented as Python classes.

What is a Python Enum?

Let’s discuss what a Python enum is. Short for enumerations, this language feature defines a set of names that are bound to constant values such as numbers, strings etc. Python enums are useful to represent data that represent a finite set of states such as days of the week, months of the year, etc.

They were added to Python 3.4 via PEP 435. However, it is available all the way back to 2.4 via pypy. As such, you can expect them to be a staple as you explore Python programming.

Illustration of woman at computer

Simple Example

We define a simple enum class (i.e. a class derived from Enum) containing the months of the year. Each month (i.e. enum member) is assigned a unique numeric constant.

from enum import Enum, unique, Flag

class Months(Enum) :
    JANUARY=1
    FEBRUARY=2
    MARCH =  3
    APRIL=4
    MAY=5
    JUNE=6
    JULY=7
    AUGUST=8
    SEPTEMBER=9
    OCTOBER=10
    NOVEMBER=11
    DECEMBER=12

Printing out the members of the Python enum

There are several ways we can do this.

# by numerical index
print (Months(7)
Months.JULY

# item index
print (Months['JULY'])
Months.JULY

# by name
print (Months.JULY)
Months.JULY

# by name
print (Months.JULY.name)
JULY

# by value
print (Months.JULY.value)
JULY

Ensuring uniqueness

Adding a @unique decorator to the class, ensures that duplicate elements don’t exist in the Python enum.

from enum import Enum, unique, Flag

>> @unique
... class Months(Enum) :
...     JANUARY=1
...     JANUARY=1
...     FEBRUARY=2
...     MARCH =  3
...     APRIL=4
...     MAY=5
...     JUNE=6
...     JULY=7
...     AUGUST=8
...     SEPTEMBER=9
...     OCTOBER=10
...     NOVEMBER=11
...     DECEMBER=12
... 
Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
  File "<stdin>", line 4, in Months
  File "/home/raghavan/anaconda3/lib/python3.8/enum.py", line 112, in __setitem__
    raise TypeError('Attempted to reuse key: %r' % key)
TypeError: Attempted to reuse key: 'JANUARY'

A TypeError occurs when the Python interpreter tries to create this Python enum.

Alternative way to create an Enum.

The Python Enum class is callable and the following Functional API can be invoked to create it.

>> quarter1= Enum('Q1', [('January', 1), ('February', 2), ('March', 3)])
>> print (quarter1(3))
Q1.March

# can also be written as
>> quarter1= Enum('Q1', ('January February March'))
>> print (quarter1)
<enum 'Q1'>

# item access
>>> quarter1['January']
<Q1.January: 1>

>>> Quarter1.March
<Q1.March: 3>

>>> Quarter1.March.value
3

Pictures of coding on a computer screen

Iterating over the elements

For loop

A simple for loop can be used to print out the members.

# simple for loop
for month in (Months) :
    print (month)

Months.JANUARY
Months.FEBRUARY
Months.MARCH
Months.APRIL
Months.MAY
Months.JUNE
Months.JULY
Months.AUGUST
Months.SEPTEMBER
Months.OCTOBER
Months.NOVEMBER
Months.DECEMBER

The __members__ attribute

__members__ is a read-only class level attribute providing a mapping from names to members. It can be iterated over to produce a similar output as above.

# iteration over elements
for name, member in Months.__members__.items():
    print (name, member)

JANUARY Months.JANUARY
FEBRUARY Months.FEBRUARY
MARCH Months.MARCH
APRIL Months.APRIL
MAY Months.MAY
JUNE Months.JUNE
JULY Months.JULY
AUGUST Months.AUGUST
SEPTEMBER Months.SEPTEMBER
OCTOBER Months.OCTOBER
NOVEMBER Months.NOVEMBER
DECEMBER Months.DECEMBER

Hashing

Python enums can be used as dictionary keys as follows.

>> months = {}
>> months[Months.JULY] = 'Many Birthdays'
>> months[Months.JANUARY] = 'First Month of the Year'

>> print (months[Months(7)])
Many Birthdays

Auto values

The values corresponding to the names can be populated automatically using auto() as demonstrated below.

# member values using auto
from enum import auto
class Quarter(Enum):
    Q1 = auto()
    Q2 = auto()
    Q3 = auto()
    Q4 = auto()


for qtr in Quarter:
    print (qtr.value)

# Output
1
2
3
4

The values, by default, are numeric. However, they can be converted into say string values by overriding _generate_next_value_ in the class.

# member values using auto
from enum import auto
class Quarter(Enum):
    def _generate_next_value_(name, start, count, last_values):
         return "[" + name + "]"
    Q1 = auto()
    Q2 = auto()
    Q3 = auto()
    Q4 = auto()

# test the values
for qtr in Quarter:
    print (qtr.value)

# Output
[Q1]
[Q2]
[Q3]
[Q4]

Woman's hands typing on a laptop

Derived Enumerations

Flag

Flag is quite similar to Enum except it has support for bitwise operations (|, &, ^, ~). These operators can be used to combine multiple Python enum elements into a mask.

The class Planets is derived from Flag and contains the 8 planets currently recognized by the International Astronomical Union. The values of the elements are required to be multiples of two while combinations need not follow that restriction.

class Planets(Flag):
    MERCURY = 1
    VENUS   = 2
    EARTH   = 4
    MARS    = 8
    SATURN  = 10
    URANUS  = 12
    NEPTUNE = 14

my_planet = {
              Planets.MERCURY: "Red Planet",
              Planets.EARTH:   "People !!!!!",
              Planets.MARS:    "Martians !!!!"
            }

# bitwise OR mask => Mercury / Earth.
mercury_earth = Planets.MERCURY|Planets.EARTH

# we check which planets in the given dictionary are
# MARS or EARTH by doing a Bitwise AND with the mask above.

for planet_name in myplanet.keys() :
    found = bool(planet_name & mercury_earth)
    print ("Found:" + str(planet_name) + ":" + str(found))

# Output

Found:Planets.MERCURY:True
Found:Planets.EARTH:True
Found:Planets.MARS:False
  1. The Planets class derives from Flag. It contains the names and values of the 8 planets.
  2. The my_planet dictionary contains the enum names of Mercury, Earth and Mars as keys.
  3. Create a bit mask by OR-ing Planets.MERCURY and Planets.EARTH.
  4. Iterate through the planet names in the dictionary and do a bit wise AND between the planet name and the mask created in Step 3. Convert the resultant value into a boolean value.
  5. The bitwise & operation will return True where the element key matches Planets.MERCURY or Planets.EARTH and False for any other planet. The output reflects this.

IntEnum

IntEnum is a derived Python enum that also derives from int. Consequently it supports int operations. It provides ordered comparisons that are not provided by the Python Enum class.

>>> from enum import Enum, auto
>>> class Quarter(Enum):
...     Q1 = auto()
...     Q2 = auto()
...     Q3 = auto()
...     Q4 = auto()


# numerical comparisons other than '==' and '!='
>>> Quarter.Q1 < Quarter.Q2
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: '<' not supported between instances of 'Quarter' and 'Quarter'


# IntEnum supports comparisons, sorting.

>>> from enum import IntEnum, Enum, auto
>>> class Quarter (IntEnum):
...     Q1 = auto() 
...     Q2 = auto() 
...     Q3 = auto() 
...     Q4 = auto()
...
>>> Quarter.Q1
<Quarter.Q1: 1>
>>> Quarter.Q2
<Quarter.Q2: 2>

>>> Quarter.Q1 < Quarter.Q2
True

>>> sorted ([Quarter.Q2, Quarter.Q1, Quarter.Q4, Quarter.Q3])
[<Quarter.Q1: 1>, <Quarter.Q2: 2>, <Quarter.Q3: 3>, , <Quarter.Q4: 4>]

# integer operations
>>> Quarter.Q1 + 10
11
  1. Inherit Quarter from Enum.
  2. Try a “<” numerical comparison between Quarter.Q1 and Quarter.Q2. It fails with a TypeError.
  3. Redefine Quarter by inheriting Quarter from IntEnum,
  4. The same comparison now succeeds.
  5. We are also able to sort the members of the IntEnum.
  6. Finally, we can perform integer operations on the IntEnum member.

IntFlag

This variation of enum is a subclass of int. The members of this Python enum are ints. These values can be combined using bitwise operators and the result is also an IntFlag member. Any other operators used on the members could remove them from IntFlag membership.

from enum import IntFlag, auto
class DaysOfWeek(IntFlag) :
    Monday = auto()
    Tuesday = auto()
    Wednesday = auto()
    Thursday = auto()
    Friday = auto()
    Saturday = auto()
    Sunday = auto()

week_end = DaysOfWeek.Saturday | DaysOfWeek.Sunday

is_mon_weekend = DaysOfWeek.Monday & week_end

print (bool(is_mon_weekend))
False

is_sun_weekend = DaysOfWeek.Sunday & week_end
print (bool(is_sun_weekend))
True
  1. We inherit DaysOfWeek from IntFlag. It contains the days of the week as members.
  2. We create a weekend mask (week_end) by bitwise-ORing Saturday and Sunday.
  3. We bitwise-AND this mask with Monday which correctly gives us a False value.
  4. We bitwise-AND this mask with Sunday which correctly gives us a True value.

Lego business man looking stressed at desk

Custom Ordering

Both IntEnum and IntFlag allows us to do custom ordering of the members where natural ordering does not help.

IntEnum

"GREEN" < "BLUE"
False

class Colors(IntEnum) :
    GREEN = 1
    BLUE  = 2

Colors.GREEN < Colors.BLUE
True
  1. The string GREEN cannot be alphabetically made less than the string BLUE.
  2. However, defining them as members of an IntEnum makes this possible.

References / Additional Reading

We’ve covered the basics, but now it’s time for you to explore Python enums yourself! Check out some of the links below to get yourself started, either with enums or with further Python programming topics.

Conclusion

Python enums are extremely useful in managing data that takes a finite set of states. They are also useful in imposing a custom ordering on data. As discussed, there are 3 variations of Python Enum i.e. IntEnum, IntFlag, and Flag. Flag and IntFlag can be used in combining states using bitwise operators, while IntEnum and IntFlag are useful for numeric comparisons and sorting.

Where you go with this newfound knowledge is up to you, though. Perhaps you’re interested in making a calendar program suited for a business suite of software. Or, maybe you’re interested in making a game and using enumerations in Python for different status effects. The uses here are pretty limitless, but we hope you’ve come away with some new information at your disposal and are able to use the Python enum to your benefit and success!

BUILD GAMES

FINAL DAYS: Unlock 250+ coding courses, guided learning paths, help from expert mentors, and more.