Python中的 enum 模块源码详析
起步
上一篇《Python的枚举类型》文末说有机会的话可以看看它的源码。那就来读一读，看看枚举的几个重要的特性是如何实现的。
要想阅读这部分，需要对元类编程有所了解。
成员名不允许重复
这部分我的第一个想法是去控制__dict__中的key。但这样的方式并不好，__dict__范围大，它包含该类的所有属性和方法。而不单单是枚举的命名空间。我在源码中发现enum使用另一个方法。通过__prepare__魔术方法可以返回一个类字典实例，在该实例使用__prepare__魔术方法自定义命名空间，在该空间内限定成员名不允许重复。
#自己实现
class_Dict(dict):
def__setitem__(self,key,value):
ifkeyinself:
raiseTypeError('Attemptedtoreusekey:%r'%key)
super().__setitem__(key,value)
classMyMeta(type):
@classmethod
def__prepare__(metacls,name,bases):
d=_Dict()
returnd
classEnum(metaclass=MyMeta):
pass
classColor(Enum):
red=1
red=1#TypeError:Attemptedtoreusekey:'red'
再看看Enum模块的具体实现：
class_EnumDict(dict):
def__init__(self):
super().__init__()
self._member_names=[]
...
def__setitem__(self,key,value):
...
elifkeyinself._member_names:
#descriptoroverwritinganenum?
raiseTypeError('Attemptedtoreusekey:%r'%key)
...
self._member_names.append(key)
super().__setitem__(key,value)
classEnumMeta(type):
@classmethod
def__prepare__(metacls,cls,bases):
enum_dict=_EnumDict()
...
returnenum_dict
classEnum(metaclass=EnumMeta):
...
模块中的_EnumDict创建了_member_names列表来存储成员名，这是因为不是所有的命名空间内的成员都是枚举的成员。比如__str__,__new__等魔术方法就不是了，所以这边的__setitem__需要做一些过滤：
def__setitem__(self,key,value):
if_is_sunder(key):#下划线开头和结尾的，如_order__
raiseValueError('_names_arereservedforfutureEnumuse')
elif_is_dunder(key):#双下划线结尾的,如__new__
ifkey=='__order__':
key='_order_'
elifkeyinself._member_names:#重复定义的key
raiseTypeError('Attemptedtoreusekey:%r'%key)
elifnot_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注入到目标类中 forname,valueincls.members.items(): member=super().__new__() member.name=name member.value=value setattr(enum_class,name,member) returnenum_class
来看下一个可运行的demo：
class_Dict(dict):
def__init__(self):
super().__init__()
self._member_names=[]
def__setitem__(self,key,value):
ifkeyinself:
raiseTypeError('Attemptedtoreusekey:%r'%key)
ifnotkey.startswith("_"):
self._member_names.append(key)
super().__setitem__(key,value)
classMyMeta(type):
@classmethod
def__prepare__(metacls,name,bases):
d=_Dict()
returnd
def__new__(metacls,cls,bases,classdict):
__new__=bases[0].__new__ifbaseselseobject.__new__
#创建枚举类
enum_class=super().__new__(metacls,cls,bases,classdict)
#创建成员
formember_nameinclassdict._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)
returnenum_class
classMyEnum(metaclass=MyMeta):
pass
classColor(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模块的代码来说明其实现了，从模块的使用特性中可以知道，如果成员值相同，后者会是前者的一个别名：
fromenumimportEnum classColor(Enum): red=1 _red=1 print(Color.redisColor._red)#True
从这可以知道，red和_red是同一对象。这又要怎么实现呢？
元类会为枚举类创建_member_map_属性来存储成员名与成员的映射关系，如果发现创建的成员的值已经在映射关系中了，就会用映射表中的对象来取代：
classEnumMeta(type): def__new__(metacls,cls,bases,classdict): ... #createournewEnumtype enum_class=super().__new__(metacls,cls,bases,classdict) enum_class._member_names_=[]#namesindefinitionorder enum_class._member_map_=OrderedDict()#name->valuemap formember_nameinclassdict._member_names: enum_member=__new__(enum_class) #Ifanothermemberwiththesamevaluewasalreadydefined,the #newmemberbecomesanaliastotheexistingone. forname,canonical_memberinenum_class._member_map_.items(): ifcanonical_member._value_==enum_member._value_: enum_member=canonical_member#取代 break else: #Aliasesdon'tappearinmembernames(onlyin__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_:
classEnumMeta(type):
def__new__(metacls,cls,bases,classdict):
...
#createournewEnumtype
enum_class=super().__new__(metacls,cls,bases,classdict)
enum_class._value2member_map_={}
formember_nameinclassdict._member_names:
value=enum_members[member_name]
enum_member=__new__(enum_class)
enum_class._value2member_map_[value]=enum_member
...
然后在Enum的__new__返回该单例即可：
classEnum(metaclass=EnumMeta):
def__new__(cls,value):
iftype(value)iscls:
returnvalue
#尝试从_value2member_map_获取
try:
ifvalueincls._value2member_map_:
returncls._value2member_map_[value]
exceptTypeError:
#从_member_map_映射获取
formemberincls._member_map_.values():
ifmember._value_==value:
returnmember
raiseValueError("%risnotavalid%s"%(value,cls.__name__))
迭代的方式遍历成员
枚举类支持迭代的方式遍历成员，按定义的顺序，如果有值重复的成员，只获取重复的第一个成员。对于重复的成员值只获取第一个成员，正好属性_member_names_只会记录第一个：
classEnum(metaclass=EnumMeta): def__iter__(cls): return(cls._member_map_[name]fornameincls._member_names_)
总结
enum模块的核心特性的实现思路就是这样，几乎都是通过元类黑魔法来实现的。对于成员之间不能做比较大小但可以做等值比较。这反而不需要讲，这其实继承自object就是这样的，不用额外做什么就有的“特性”了。
总之，enum模块相对独立，且代码量不多，对于想知道元类编程可以阅读一下，教科书式教学，还有单例模式等，值得一读。
好了，以上就是这篇文章的全部内容了，希望本文的内容对大家的学习或者工作具有一定的参考学习价值，如果有疑问大家可以留言交流，谢谢大家对毛票票的支持。