欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页  >  IT编程

Python中的 enum 模块源码详析

程序员文章站 2023-10-20 11:47:03
起步 上一篇 《python 的枚举类型》 文末说有机会的话可以看看它的源码。那就来读一读,看看枚举的几个重要的特性是如何实现的。 要想阅读这部分,需要对元类编程有...

起步

上一篇 《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 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 模块相对独立,且代码量不多,对于想知道元类编程可以阅读一下,教科书式教学,还有单例模式等,值得一读。

好了,以上就是这篇文章的全部内容了,希望本文的内容对大家的学习或者工作具有一定的参考学习价值,如果有疑问大家可以留言交流,谢谢大家对的支持。