Python的enum 模块源码分析

作者丨weapon

zhihu.com/people/hong-wei-peng

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

成员名不允许重复

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

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

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

模块考虑的会更全面。

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

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

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

来看下一个可运行的demo：

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

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

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

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

元类会为枚举类创建_member_map_ 属性来存储成员名与成员的映射关系，如果发现创建的成员的值已经在映射关系中了，就会用映射表中的对象来取代：

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

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

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

枚举类中的成员都是单例模式，元类创建的枚举类中还维护了值到成员的映射关系 _value2member_map_:

然后在 Enum 的__new__返回该单例即可：

迭代的方式遍历成员

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

总结

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

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

推荐↓↓↓