Detailed analysis of the source code of the enum module in Python (code example)

This article brings you a detailed analysis (code example) of the enum module source code in Python. It has certain reference value. Friends in need can refer to it. I hope it will be helpful to you.

Previous article "Detailed explanation of enumeration types in Python (code examples)" At the end of the article, it is said that if you have the opportunity, you can take a look at its source code. Then read it and see how several important features of enumeration are implemented.

To read this part, you need to have some understanding of metaclass programming.

Repeated member names are not allowed

My first idea for this part is to control the key in __dict__. But this way is not good, __dict__ has a large scope, it contains all attributes and methods of the class. Not just the enumeration namespace. I found in the source code that enum uses another method. A dictionary-like instance can be returned through the __prepare__ magic method. In this instance, the __prepare__ magic method is used to customize the namespace, and repeated member names are not allowed in this space.

# 自己实现
class _Dict(dict):
    def __setitem__(self, key, value):
        if key in self:
            raise TypeError('Attempted to reuse key: %r' % key)
        super().__setitem__(key, value)

class MyMeta(type):
    @classmethod
    def __prepare__(metacls, name, bases):
        d = _Dict()
        return d

class Enum(metaclass=MyMeta):
    pass

class Color(Enum):
    red = 1
    red = 1         # TypeError: Attempted to reuse key: 'red'

Let’s look at the specific implementation of the Enum module:

class _EnumDict(dict):
    def __init__(self):
        super().__init__()
        self._member_names = []
        ...

    def __setitem__(self, key, value):
        ...
        elif key in self._member_names:
            # descriptor overwriting an enum?
            raise TypeError('Attempted to reuse key: %r' % key)
        ...
        self._member_names.append(key)
        super().__setitem__(key, value)
        
class EnumMeta(type):
    @classmethod
    def __prepare__(metacls, cls, bases):
        enum_dict = _EnumDict()
        ...
        return enum_dict

class Enum(metaclass=EnumMeta):
    ...

The _EnumDict in the module creates a _member_names list to store member names. This is because not all members in the namespace are Is a member of the enumeration. For example, __str__, __new__ and other magic methods are not, so the __setitem__ here needs to do some filtering:

def __setitem__(self, key, value):
    if _is_sunder(key):     # 下划线开头和结尾的，如 _order__
        raise ValueError('_names_ are reserved for future Enum use')
    elif _is_dunder(key):   # 双下划线结尾的, 如 __new__
        if key == '__order__':
            key = '_order_'
    elif key in self._member_names: # 重复定义的 key
        raise TypeError('Attempted to reuse key: %r' % key)
    elif not _is_descriptor(value): # value得不是描述符
        self._member_names.append(key)
        self._last_values.append(value)
    super().__setitem__(key, value)

The module will consider it more comprehensively.

Each member has a name attribute and a value attribute

In the above code, the value obtained by Color.red is 1. In the eumu module, each member of the defined enumeration class has a name and attribute value; and if you are careful, you will find that Color.red is an example of Color. How is this situation achieved?

It is still done with metaclasses and implemented in __new__ of metaclasses. The specific idea is to first create the target class, then create the same class for each member, and then use setattr to add subsequent The class is added to the target class as an attribute. The pseudo code is as follows:

def __new__(metacls, cls, bases, classdict):
    __new__ = cls.__new__
    # 创建枚举类
    enum_class = super().__new__()
    # 每个成员都是cls的示例，通过setattr注入到目标类中
    for name, value in cls.members.items():
        member = super().__new__()
        member.name = name
        member.value = value
        setattr(enum_class, name, member)
    return enum_class

Let’s look at the next runnable demo:

class _Dict(dict):
    def __init__(self):
        super().__init__()
        self._member_names = []

    def __setitem__(self, key, value):
        if key in self:
            raise TypeError('Attempted to reuse key: %r' % key)

        if not key.startswith("_"):
            self._member_names.append(key)
        super().__setitem__(key, value)

class MyMeta(type):
    @classmethod
    def __prepare__(metacls, name, bases):
        d = _Dict()
        return d

    def __new__(metacls, cls, bases, classdict):
        __new__ = bases[0].__new__ if bases else object.__new__
        # 创建枚举类
        enum_class = super().__new__(metacls, cls, bases, classdict)

        # 创建成员
        for member_name in classdict._member_names:
            value = classdict[member_name]
            enum_member = __new__(enum_class)
            enum_member.name = member_name
            enum_member.value = value
            setattr(enum_class, member_name, enum_member)

        return enum_class

class MyEnum(metaclass=MyMeta):
    pass

class Color(MyEnum):
    red = 1
    blue = 2

    def __str__(self):
        return "%s.%s" % (self.__class__.__name__, self.name)

print(Color.red)        # Color.red
print(Color.red.name)   # red
print(Color.red.value)  # 1

The enum module allows each member to have a name and value. The implementation idea of ​​attributes is the same (I won’t post the code). EnumMeta.__new__ is the focus of this module, and almost all enumeration features are implemented in this function.

When the member values ​​are the same, the second member is an alias of the first member

From this section onwards, the description of the class implemented by yourself will no longer be used. , but illustrates its implementation by disassembling the code of the enum module. From the usage characteristics of the module, we can know that if the member values ​​​​are the same, the latter will be an alias of the former:

from enum import Enum
class Color(Enum):
    red = 1
    _red = 1

print(Color.red is Color._red)  # True

You can know from this , red and _red are the same object. How to achieve this?

The metaclass will create the _member_map_ attribute for the enumeration class to store the mapping relationship between member names and members. If it is found that the value of the created member is already in the mapping relationship, the object in the mapping table will be used. Replace:

class EnumMeta(type):
    def __new__(metacls, cls, bases, classdict):
        ...
        # create our new Enum type
        enum_class = super().__new__(metacls, cls, bases, classdict)
        enum_class._member_names_ = []               # names in definition order
        enum_class._member_map_ = OrderedDict()      # name->value map

        for member_name in classdict._member_names:
            enum_member = __new__(enum_class)

            # If another member with the same value was already defined, the
            # new member becomes an alias to the existing one.
            for name, canonical_member in enum_class._member_map_.items():
                if canonical_member._value_ == enum_member._value_:
                    enum_member = canonical_member     # 取代
                    break
            else:
                # Aliases don't appear in member names (only in __members__).
                enum_class._member_names_.append(member_name)  # 新成员，添加到_member_names_中
            
            enum_class._member_map_[member_name] = enum_member
            ...

From a code point of view, even if the member values ​​are the same, objects will still be created for them all first, but the objects created later will be garbage collected soon (I think there is optimization here spatial). By comparing with the _member_map_ mapping table, the member used to create the member value replaces the subsequent ones, but both member names will be in the _member_map_. For example, red and _red in the example are both in the dictionary, but they point to the same an object.

The attribute _member_names_ will only record the first one, which will be related to the iteration of the enumeration.

Members can be obtained through member values

print(Color['red'])  # Color.red  通过成员名来获取成员
print(Color(1))      # Color.red  通过成员值来获取成员

The members in the enumeration class are singleton mode, and the values ​​​​are also maintained in the enumeration class created by the metaclass Mapping relationship to members_value2member_map_:

class EnumMeta(type):
    def __new__(metacls, cls, bases, classdict):
        ...
        # create our new Enum type
        enum_class = super().__new__(metacls, cls, bases, classdict)
        enum_class._value2member_map_ = {}

        for member_name in classdict._member_names:
            value = enum_members[member_name]
            enum_member = __new__(enum_class)

            enum_class._value2member_map_[value] = enum_member
            ...

Then return the singleton in Enum’s __new__:

class Enum(metaclass=EnumMeta):
    def __new__(cls, value):
        if type(value) is cls:
            return value

        # 尝试从 _value2member_map_ 获取
        try:
            if value in cls._value2member_map_:
                return cls._value2member_map_[value]
        except TypeError:
            # 从 _member_map_ 映射获取
            for member in cls._member_map_.values():
                if member._value_ == value:
                    return member

        raise ValueError("%r is not a valid %s" % (value, cls.__name__))

Iteratively traverse the members

The enumeration class supports iterative traversal of members in the defined order. If there are members with duplicate values, only the first duplicate member is obtained. For duplicate member values, only the first member is obtained, and the attribute _member_names_ will only record the first one:

class Enum(metaclass=EnumMeta):
    def __iter__(cls):
        return (cls._member_map_[name] for name in cls._member_names_)

Summary

Implementation of the core features of the enum module This is the idea, almost all achieved through meta-class black magic. The members cannot be compared in size but can be compared in equal value. There is no need to talk about this. This is actually inherited from object. It has "features" without doing anything extra.

The above is the detailed content of Detailed analysis of the source code of the enum module in Python (code example). For more information, please follow other related articles on the PHP Chinese website!