Python中enum模块源码的详细分析（代码示例）

百变鹏仔 1个月前 (01-23) #Python

文章标签示例

本篇文章给大家带来的内容是关于python中enum模块源码的详细分析（代码示例），有一定的参考价值，有需要的朋友可以参考一下，希望对你有所帮助。

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

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

成员名不允许重复

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

立即学习“Python免费学习笔记（深入）”；

# 自己实现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 的方式将后续的类作为属性添加到目标类中，伪代码如下：

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_classclass MyEnum(metaclass=MyMeta):    passclass Color(MyEnum):    red = 1    blue = 2    def __str__(self):        return "%s.%s" % (self.__class__.__name__, self.name)print(Color.red)        # Color.redprint(Color.red.name)   # redprint(Color.red.value)  # 1

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

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

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

from enum import Enumclass Color(Enum):    red = 1    _red = 1print(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 就是这样的，不用额外做什么就有的“特性”了。

文章推荐

Python中enum模块源码的详细分析（代码示例）

Python实现字典的key和values的交换

使用Python脚本来获取Cisco设备信息的示例

Python的Django中django-userena组件的简单使用教程

零基础写python爬虫之神器正则表达式

零基础写python爬虫之抓取百度贴吧代码分享