Python 的 enum 模块源码分析

Posted Python中文社区

tags:

篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了Python 的 enum 模块源码分析相关的知识，希望对你有一定的参考价值。

♚

博客：zhihu.com/people/hong-wei-peng

要想阅读这部分，需要对元类编程有所了解。

成员名不允许重复

这部分我的第一个想法是去控制__dict__中的 key 。但这样的方式并不好，__dict__ 范围大，它包含该类的所有属性和方法。而不单单是枚举的命名空间。我在源码中发现 enum 使用另一个方法。通过 __prepare__ 魔术方法可以返回一个类字典实例，在该实例 使用__prepare__ 魔术方法自定义命名空间，在该空间内限定成员名不允许重复。

# 自己实现
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'

再看看 Enum 模块的具体实现：

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):
    ...

模块中的_EnumDict 创建了_member_names 列表来存储成员名，这是因为不是所有的命名空间内的成员都是枚举的成员。比如 __str__, __new__ 等魔术方法就不是了，所以这边的 __setitem__ 需要做一些过滤：

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)

模块考虑的会更全面。

每个成员都有名称属性和值属性

上述的代码中，Color.red 取得的值是 1。而 eumu 模块中，定义的枚举类中，每个成员都是有名称和属性值的；并且细心的话还会发现 Color.red 是 Color 的示例。这样的情况是如何来实现的呢。

还是用元类来完成，在元类的 __new__ 中实现，具体的思路是，先创建目标类，然后为每个成员都创建一样的类，再通过 setattr 的方式将后续的类作为属性添加到目标类中，伪代码如下：

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

来看下一个可运行的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

enum 模块在让每个成员都有名称和值的属性的实现思路是一样的（代码我就不贴了）。EnumMeta.__new__ 是该模块的重点，几乎所有枚举的特性都在这个函数实现。

当成员值相同时，第二个成员是第一个成员的别名

从这节开始就不再使用自己实现的类的说明了，而是通过拆解 enum 模块的代码来说明其实现了，从模块的使用特性中可以知道，如果成员值相同，后者会是前者的一个别名：

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

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

从这可以知道，red和_red是同一对象。这又要怎么实现呢？

元类会为枚举类创建_member_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._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
            ...

从代码上来看，即使是成员值相同，还是会先为他们都创建对象，不过后创建的很快就会被垃圾回收掉了(我认为这边是有优化空间的)。通过与_member_map_ 映射表做对比，用以创建该成员值的成员取代后续，但两者成员名都会在 _member_map_中，如例子中的red和_red 都在该字典，但他们指向的是同一个对象。

属性_member_names_只会记录第一个，这将会与枚举的迭代有关。

可以通过成员值来获取成员

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

枚举类中的成员都是单例模式，元类创建的枚举类中还维护了值到成员的映射关系 _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
            ...

然后在 Enum 的__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__))

迭代的方式遍历成员

枚举类支持迭代的方式遍历成员，按定义的顺序，如果有值重复的成员，只获取重复的第一个成员。对于重复的成员值只获取第一个成员，正好属性_member_names_只会记录第一个：

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

总结

enum 模块的核心特性的实现思路就是这样，几乎都是通过元类黑魔法来实现的。对于成员之间不能做比较大小但可以做等值比较。这反而不需要讲，这其实继承自 object 就是这样的，不用额外做什么就有的“特性”了。

总之，enum 模块相对独立，且代码量不多，对于想知道元类编程可以阅读一下，教科书式教学，还有单例模式等，值得一读。

Python 的 enum 模块源码分析

《数据分析实战45讲》


作者陈旸，清华大学计算机系博士，10岁开始编程，2次获得全国信息学奥林匹克竞赛一等奖，2次ACM国际编程比赛亚洲区铜奖。并通过数据分析，先后为腾讯视频、易车、58同城、蚂蚁金服、京东分析用户画像和传播话题，为品牌活动做传播决策。


Python 的 enum 模块源码分析


还有好福利，极客时间2019年限量版周历（价值¥99），我们会从今天扫海报购买粉丝中，随机免费送给3人（今天24:00前截止！）

   

     

§ §

Python中文社区作为一个去中心化的全球技术社区，以成为全球20万Python中文开发者的精神部落为愿景，目前覆盖各大主流媒体和协作平台，与阿里、腾讯、百度、微软、亚马逊、开源中国、CSDN等业界知名公司和技术社区建立了广泛的联系，拥有来自十多个国家和地区数万名登记会员，会员来自以公安部、工信部、清华大学、北京大学、北京邮电大学、中国人民银行、中科院、中金、华为、BAT、谷歌、微软等为代表的政府机关、科研单位、金融机构以及海内外知名公司，全平台近20万开发者关注。

最近热门











投稿邮箱：pythonpost@163.com


▼ 点击下方阅读原文，免费成为社区注册会员 

以上是关于Python 的 enum 模块源码分析的主要内容，如果未能解决你的问题，请参考以下文章

《Docker 源码分析》全球首发啦！

Python 枚举

Android 插件化VirtualApp 源码分析 ( 目前的 API 现状 | 安装应用源码分析 | 安装按钮执行的操作 | 返回到 HomeActivity 执行的操作 )(代码片段

基于Python的datetime模块和time模块源码阅读分析

Android 逆向整体加固脱壳 ( DEX 优化流程分析 | DexPrepare.cpp 中 dvmOptimizeDexFile() 方法分析 | /bin/dexopt 源码分析 )(代码片段

python3 enum模块