Enums, short for enumerations, are a powerful feature in Python that often goes underutilized. They provide a way to define symbolic names bound to unique, constant values. In this blog post, we will explore the various enum classes in Python and demonstrate how they can be mixed with other classes to extend their functionality, enhancing your Python code in elegant and efficient ways.

What is an enum?

An enum is a symbolic name for a set of values. Enums are used when we have a fixed set of related constants that make the code more readable and less error-prone.

Let’s start with a basic example:

from enum import Enum

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

print(Color.RED)
print(Color.GREEN.value)
print(Color(2))
print(Color["RED"])

Output:

Color.RED
2
Color.GREEN
Color.RED

Here, Color is an enumeration with three members: RED, GREEN, and BLUE. Each member has a unique value, and if you define two members with the same value, the second will be an alias of the first.

As you can see, enums provide a number of ways to access the attributes.

Color.GREEN : Using dot notation will yield the member itself, and to access the value; you can then call .value

Color(2) : This is a lookup by value; in the case of aliased values, will return the first member with the given value.

Color["RED"] : This performs a key lookup, finding the member with the given name.

What are the benefits of using enums?

Enums offer several advantages that make your code more robust and maintainable:

1. Improved readability: Enums provide meaningful names to constant values, making the code easier to understand.
2. Reduced errors: By using Enums, you avoid the pitfalls of using literal constants, reducing the chances of typos and invalid values.
3. Namespace Organization: Enums keep related constants grouped together, avoiding clutter in the global namespace.

Different types of enums

Python’s enum module provides several different classes to create enumerations:

1. Enum: The base class for creating simple enumerations.
2. IntEnum: A subclass of Enum where members are also integers and can be used in place of integers.
3. Flag: A subclass of Enum where members can be combined using bitwise operations.
4. IntFlag: A combination of IntEnum and Flag.

Using IntEnum

IntEnum can be useful when you need enum members to behave like integers:

from enum import IntEnum

class Status(IntEnum):
    SUCCESS = 1
    FAILURE = 2
    PENDING = 3

def check_status(status):
    if status == Status.SUCCESS:
        print("Operation was successful!")

check_status(1)  # Works because Status.SUCCESS is also an integer

Output:

Operation was successful!

Using Flag and IntFlag

Flags are useful when you need to represent combinations of values:

from enum import Flag, auto

class Permission(Flag):
    READ = auto()
    WRITE = auto()
    EXECUTE = auto()

# Combining flags
perms = Permission.READ | Permission.WRITE
print(perms)
print(Permission.READ in perms)

Output:

Permission.READ|WRITE
True

We can use the auto object to allow the Enum class to automatically generate the numerical value for the enum’s member.

Extending enum functionality

Enums on their own provide some useful features, but you can also extend this functionality to suit your needs.

Adding methods to enums

Enums can have methods just like regular classes:

from enum import Enum

class Weekday(Enum):
    MONDAY = 1
    TUESDAY = 2
    WEDNESDAY = 3
    THURSDAY = 4
    FRIDAY = 5
    SATURDAY = 6
    SUNDAY = 7

    def is_weekend(self):
        return self in (Weekday.SATURDAY, Weekday.SUNDAY)

print(Weekday.SATURDAY.is_weekend())
print(Weekday.WEDNESDAY.is_weekend())

Output:

True
False

Custom values

Adding custom values to an enum can come with a gotcha!

class Color(Enum):
    RED = 1, (255, 0, 0)
    GREEN = 2, (0, 255, 0)
    BLUE = 3, (0, 0, 255)

    def __init__(self, value, rgb):
        self.value = value  # Attribute Error: Cannot set `value`
        self.rgb = rgb

The value attribute is a protected attribute. To set it, we instead need to set the _value_ attribute.

class Color(Enum):
    RED = 1, (255, 0, 0)
    GREEN = 2, (0, 255, 0)
    BLUE = 3, (0, 0, 255)

    def __init__(self, value, rgb):
        self._value_ = value
        self.rgb = rgb

print(Color.RED)
print(Color.RED.value)
print(Color.RED.rgb)

Output:

Color.RED
1
(255, 0, 0)

This approach allows us to include additional information alongside the main value of the enum members.

Final thoughts

Enums in Python are a powerful tool for defining and managing sets of constants. By using different types of enums like Enum, IntEnum, Flag, and IntFlag, readability is improved, errors are minimized, and flexibility is increased. Try using enums for managing constants or designing strategies to simplify code, improve expressiveness, and receive more helpful hints from your IDE.