Python中的 enum 模块源码详析

起步

上一篇 《Python 的枚举类型》 文末说有机会的话可以看看它的源码。那就来读一读，看看枚举的几个重要的特性是如何实现的。

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

成员名不允许重复

这部分我的第一个想法是去控制 __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 dclass Enum(metaclass=MyMeta): passclass 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_dictclass 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 的方式将后续的类作为属性添加到目标类中，伪代码如下：