Introduction

Why study the implementation of Python’s enum module? Python is an object-oriented programming (OOP) language, and the enum module in the standard library involves many OOP concepts, such as mixins, duck typing, metaclasses, and magic methods. Reading the enum module’s source code is an effective way to understand object-oriented programming. This article is based on the CPython 3.11 branch, corresponding to Python version 3.11.4.

Topics covered:

  • Enum classes
  • Duck typing
  • Magic methods
  • Protocols
  • Decorators
  • Descriptor classes
  • Metaclasses

Concepts

Enum Classes

The members of an enum class are special instances of the enum class itself. In other words, the attribute names remain unchanged, but their values become special instances of the enum class. In an enum class, the originally defined attribute names and values are stored in the name and value attributes of these special instances.

from enum import Enum

class Role(Enum):
    ADMIN = 'admin'
    DEVELOPER = 'developer'
    GUEST = 'guest'

print(type(Role.ADMIN))
# Output: <enum 'Role'>
# The original attribute has been replaced with an instance of Role

print(Role.ADMIN.name) # 'ADMIN'
print(Role.ADMIN.value) # 'admin'
# The original attribute name and value are stored in
# the special instance's name and value attributes

Usage of Enum Classes

In Python web development, user permission management and checking are commonly implemented using enum classes.

Different permissions correspond to different enum class instances, such as GUEST, DEVELOPER, and ADMIN in the example above.

When defining a user table with an ORM (e.g., SQLAlchemy), the user role column is set to an enum-related database type. When writing to the database, the ORM stores the enum instance’s value:

from sqlalchemy import Column, Integer, String, Enum
from sqlalchemy.ext.declarative import declarative_base
from enum import Enum as PyEnum

class Role(PyEnum):
    ADMIN = 'admin'
    DEVELOPER = 'developer'
    GUEST = 'guest'

Base = declarative_base()

class User(Base):
    __tablename__ = 'users'
    id = Column(Integer, primary_key=True)
    name = Column(String)
    role = Column(Enum(Role))

When querying the database, the ORM returns the enum instance corresponding to the stored value:

# Suppose the database query returns a user whose role is: 'admin',
# stored in the value variable
value = 'admin'
# Role(value) is actually a call to the enum class's __call__() method
# If the given enum value has a corresponding enum instance, it returns that instance
current_user.role = Role(value)

This enables object-oriented permission checking logic:

if current_user.role is Role.admin:
    # business logic

Duck Typing / Magic Methods / Protocols

Duck typing: If it walks like a duck and quacks like a duck, then it is a duck.

Magic methods: Internal methods of the form __iter__, __next__, __getitem__, __call__, __lt__, __set_name__, etc.

Protocols: Protocols are associated with magic methods. If a class implements one or more magic methods required by a protocol, it is considered to implement that protocol. Common protocols include: iterable protocol, callable protocol, context management protocol, etc.

The relationship between the three:

  • A protocol is associated with one or more magic methods
  • If a class implements one or more magic methods defined by a protocol, it implements that protocol
  • Implementing the relevant protocol (looks like a duck) allows usage in specific functions or syntax (is a duck). For example:
from collections import defaultdict

class AnimalCount(object):
    def __init__(self, **kwargs):
        self._dict = defaultdict(int)
        for key, value in kwargs.items():
            self._dict[key] = value

    def __getitem__(self, key):
        return self._dict[key]

    def __call__(self, name, new_value):
        self._dict[name] = new_value


# Create an instance of the class
my_pets = AnimalCount(cat=1, dog=2)

# Because __getitem__() is implemented, the instance can be used like a dictionary
print(my_pets['cat'])
print(my_pets['dog'])
print(my_pets['bear'])

# Because __call__() is implemented, the instance can be called like a function
my_pets('dog', 3)
print(my_pets['dog'])

Decorators

What is a decorator?

A function or class used to enhance the functionality of an existing function.

What is a decorator class?

A class that implements the __call__() magic method. An instance of the class replaces the decorated function and provides function-like calling behavior through __call__(). For example:

from functools import wraps
class MyDecorator:
    def __init__(self, arg1, arg2):
        # arg1 and arg2 here are decorator arguments
        self.arg1 = arg1
        self.arg2 = arg2

    def __call__(self, f):
        # We actually replace my_function with this wrapped function
        # wrapped does not have some of my_function's internal attributes,
        # such as: __doc__, __name__
        # The wraps decorator gives wrapped the same internal attributes
        # as my_function
        @wraps(f)
        # f here is the decorated function
        def wrapped(*args, **kwargs):
            print(f"Decorator arguments: {self.arg1}, {self.arg2}")
            f(*args, **kwargs)
        return wrapped

# Now we use the decorator and pass in arguments
@MyDecorator("hello", "world")
def my_function(a, b):
    print(f"Function arguments: {a}, {b}")

my_function("foo", "bar")

The mechanism: after the decorator class creates an instance, it calls __call__() which returns a function that replaces the original my_function.

What is a decorator function?

from functools import wraps

def repeat(num_times):
    def decorator_repeat(func):
        @wraps(func)
        def wrapper_repeat(*args, **kwargs):
            for _ in range(num_times):
                result = func(*args, **kwargs)
            return result
        # Return a new function that calls the decorated function num_times times
        return wrapper_repeat
    # Return a decorator with the argument applied;
    # the decorator's parameter is the decorated function
    return decorator_repeat

@repeat(3)
def greet(name):
    print(f"Hello {name}")

greet("World")

The mechanism is similar: pass in arguments to return a parameterized decorator, then pass in the decorated function to return a new function. The new function enhances the original function’s functionality and is ready for use.

Descriptor Classes

The most common descriptor class implementation is the property decorator. property is implemented in C. A similar Python implementation can be found at: DynamicClassAttribute

class property(object):

    def __init__(self, fget=None, fset=None, fdel=None, doc=None):
        if doc is None and fget is not None and hasattr(fget, "__doc__"):
            doc = fget.__doc__
        self.__get = fget
        self.__set = fset
        self.__del = fdel
        try:
            self.__doc__ = doc
        except AttributeError:  # read-only or dict-less class
            pass

    def __get__(self, inst, type=None):
        if inst is None:
            return self
        if self.__get is None:
            raise AttributeError, "property has no getter"
        return self.__get(inst)

    def __set__(self, inst, value):
        if self.__set is None:
            raise AttributeError, "property has no setter"
        return self.__set(inst, value)

    def __delete__(self, inst):
        if self.__del is None:
            raise AttributeError, "property has no deleter"
        return self.__del(inst)

In Python, a descriptor is a class that implements the descriptor protocol. The descriptor protocol includes three special methods: __get__(), __set__(), and __delete__(). In Python 3.6, a new method was added: __set_name__, which is also a magic method used by the enum module (more on this later). Descriptor classes allow the process of creating an attribute (reading, writing, deleting) to be split into multiple functions. Common use cases include:

  • Adding validation logic before reading or recording a value
  • Defining read-only property methods, such as the common usage of the property decorator
class Circle:
    def __init__(self, radius):
        self._radius = radius

    @property
    def radius(self):
        return self._radius

    @property
    def diameter(self):
        return self._radius * 2

c = Circle(5)
print(c.radius)  # Output: 5
print(c.diameter)  # Output: 10

Metaclasses

A metaclass is a class that creates classes. By defining a metaclass, you can control the class creation process.

class EnumType(type):
    def __prepare__(metacls, cls, bases, **kwds):
        # Will return the class dictionary

    def __new__(metacls, cls, bases, classdict, *, boundary=None, _simple=False, **kwds):
        # The actual logic for creating the class

# By specifying the metaclass parameter, the class is created based on the metaclass
class Enum(metaclass=EnumType):
    pass

Execution Order of Metaclass Magic Methods

  1. __prepare__(): This is the first method called by a metaclass when creating a class. It returns a dictionary object, also known as the class dictionary (classdict), used to store the attributes and methods of the newly created class. The attributes and methods we define in the class are recorded internally through the dictionary object’s __setitem__()
  2. __new__(): Where the class instance is actually created

How Enum Class Instances Are Created

Since enum class attributes are instances of the enum class, what actually happens in the metaclass’s __new__() is: an enum class instance is created using the attribute name as the instance name. This instance does not retain the class attributes as originally defined (because all class attributes will be converted into special instances). The attribute names and values are added to the instance’s name and value attributes:

from enum import Enum

class Role(Enum):
    ADMIN = 'admin'
    DEVELOPER = 'developer'
    GUEST = 'guest'

print(type(Role.ADMIN)) # Output: <enum 'Role'>
print(Role.ADMIN.name) # Output: ADMIN
print(Role.ADMIN.value) # Output: admin
print(Role.DEVELOPER.value) # Output: developer
print(Role.GUEST.value) # Output: guest

The creation process works as follows: first, __prepare__() prepares and returns the class dictionary. Then __new__() creates the class based on the class dictionary. When calling the parent class type’s __new__(), if the class dictionary contains instances of descriptor classes, their __set_name__() methods are called based on the duck typing principle.

Below is the descriptor class definition. Full code at: _proto_member descriptor class

class _proto_member:
    def __init__(self, value):
        self.value = value

    def __set_name__(self, enum_class, member_name):
        # content omitted

Below is the operation that adds descriptor class instances to the class dictionary:

for name in member_names:
   value = classdict[name]
   classdict[name] = _proto_member(value)

The enum module leverages this mechanism to replace class attribute values with enum class instances.

Therefore, the operation that ultimately replaces attribute values with enum class instances occurs inside the descriptor class’s __set_name__() magic method. See the code: setattr

Summary

If the operations industry demands breadth of vision (familiarity with various tools, open-source projects, and solutions), then development demands depth in a specific direction (or thoroughness – full consideration and implementation of all possibilities).

This is a record of my experience reading the Python standard library’s enum module. If you also want to understand object-oriented programming, studying the enum module’s implementation is a great idea.