Using Enumeration Classes

When we need to define constants, one approach is to use uppercase variables defined with integers—for example, for months:

JAN = 1
FEB = 2
MAR = 3
...
NOV = 11
DEC = 12

The advantage of this method is its simplicity, while the disadvantages are that the type is int and the values remain mutable variables.

A better approach is to define a class type for such enumeration types, where each constant is a unique instance of the class. Python provides the Enum class to implement this functionality:

from enum import Enum

Month = Enum('Month', ('Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec'))

This gives us an enumeration class of type Month. We can directly reference a constant using Month.Jan, or iterate over all its members:

for name, member in Month.__members__.items():
    print(name, '=>', member, ',', member.value)

The value attribute is an integer constant automatically assigned to each member, starting from 1 by default.

For more precise control over enumeration types, we can derive a custom class from Enum:

from enum import Enum, unique

@unique
class Weekday(Enum):
    Sun = 0  # Sun's value is set to 0
    Mon = 1
    Tue = 2
    Wed = 3
    Thu = 4
    Fri = 5
    Sat = 6

The @unique decorator helps us check and ensure there are no duplicate values.

There are several ways to access these enumeration types:

>>> day1 = Weekday.Mon
>>> print(day1)
Weekday.Mon
>>> print(Weekday.Tue)
Weekday.Tue
>>> print(Weekday['Tue'])
Weekday.Tue
>>> print(Weekday.Tue.value)
2
>>> print(day1 == Weekday.Mon)
True
>>> print(day1 == Weekday.Tue)
False
>>> print(Weekday(1))
Weekday.Mon
>>> print(day1 == Weekday(1))
True
>>> Weekday(7)
Traceback (most recent call last):
  ...
ValueError: 7 is not a valid Weekday
>>> for name, member in Weekday.__members__.items():
...     print(name, '=>', member)
...
Sun => Weekday.Sun
Mon => Weekday.Mon
Tue => Weekday.Tue
Wed => Weekday.Wed
Thu => Weekday.Thu
Fri => Weekday.Fri
Sat => Weekday.Sat

As shown, we can reference enumeration constants either by member name or directly by their value.

Exercise

Refactor the gender attribute of Student to an enumeration type to avoid using strings:

from enum import Enum, unique

class Gender(Enum):
    Male = 0
    Female = 1

class Student(object):
    def __init__(self, name, gender):
        self.name = name
        self.gender = gender

# Test:
bart = Student('Bart', Gender.Male)
if bart.gender == Gender.Male:
    print('Test passed!')
else:
    print('Test failed!')

Summary

Enum allows us to define a group of related constants within a class. The class is immutable, and its members can be directly compared.

Reference Source Code

use_enum.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from enum import Enum, unique


@unique
class Weekday(Enum):
    Sun = 0
    Mon = 1
    Tue = 2
    Wed = 3
    Thu = 4
    Fri = 5
    Sat = 6


day1 = Weekday.Mon

print("day1 =", day1)
print("Weekday.Tue =", Weekday.Tue)
print("Weekday['Tue'] =", Weekday["Tue"])
print("Weekday.Tue.value =", Weekday.Tue.value)
print("day1 == Weekday.Mon ?", day1 == Weekday.Mon)
print("day1 == Weekday.Tue ?", day1 == Weekday.Tue)
print("day1 == Weekday(1) ?", day1 == Weekday(1))

for name, member in Weekday.__members__.items():
    print(name, "=>", member)

Month = Enum("Month", ("Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"))

for name, member in Month.__members__.items():
    print(name, "=>", member, ",", member.value)