Python中的 enum 模块源码详析

# 自己实现
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'

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):
 ...

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)

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

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

from enum import Enum
class Color(Enum):
 red = 1
 _red = 1

print(Color.red is Color._red) # True

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
  ...

print(Color['red']) # Color.red 通过成员名来获取成员
print(Color(1)) # Color.red 通过成员值来获取成员

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
  ...

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__))

class Enum(metaclass=EnumMeta):
 def __iter__(cls):
 return (cls._member_map_[name] for name in cls._member_names_)