diff options
Diffstat (limited to 'Lib/enum.py')
| -rw-r--r-- | Lib/enum.py | 260 |
1 files changed, 246 insertions, 14 deletions
diff --git a/Lib/enum.py b/Lib/enum.py index b8787d19b88..e89c17d583e 100644 --- a/Lib/enum.py +++ b/Lib/enum.py @@ -1,8 +1,16 @@ import sys -from collections import OrderedDict from types import MappingProxyType, DynamicClassAttribute +from functools import reduce +from operator import or_ as _or_ -__all__ = ['Enum', 'IntEnum', 'unique'] +# try _collections first to reduce startup cost +try: + from _collections import OrderedDict +except ImportError: + from collections import OrderedDict + + +__all__ = ['EnumMeta', 'Enum', 'IntEnum', 'Flags', 'IntFlags', 'unique'] def _is_descriptor(obj): @@ -58,16 +66,21 @@ class _EnumDict(dict): """ if _is_sunder(key): - raise ValueError('_names_ are reserved for future Enum use') + if key not in ( + '_order_', '_create_pseudo_member_', '_decompose_', + '_generate_next_value_', '_missing_', + ): + raise ValueError('_names_ are reserved for future Enum use') elif _is_dunder(key): - pass + if key == '__order__': + key = '_order_' elif key in self._member_names: # descriptor overwriting an enum? raise TypeError('Attempted to reuse key: %r' % key) elif not _is_descriptor(value): if key in self: # enum overwriting a descriptor? - raise TypeError('Key already defined as: %r' % self[key]) + raise TypeError('%r already defined as: %r' % (key, self[key])) self._member_names.append(key) super().__setitem__(key, value) @@ -83,9 +96,15 @@ class EnumMeta(type): """Metaclass for Enum""" @classmethod def __prepare__(metacls, cls, bases): - return _EnumDict() + # create the namespace dict + enum_dict = _EnumDict() + # inherit previous flags and _generate_next_value_ function + member_type, first_enum = metacls._get_mixins_(bases) + if first_enum is not None: + enum_dict['_generate_next_value_'] = getattr(first_enum, '_generate_next_value_', None) + return enum_dict - def __new__(metacls, cls, bases, classdict): + def __new__(metacls, cls, bases, classdict, **kwds): # an Enum class is final once enumeration items have been defined; it # cannot be mixed with other types (int, float, etc.) if it has an # inherited __new__ unless a new __new__ is defined (or the resulting @@ -96,12 +115,15 @@ class EnumMeta(type): # save enum items into separate mapping so they don't get baked into # the new class - members = {k: classdict[k] for k in classdict._member_names} + enum_members = {k: classdict[k] for k in classdict._member_names} for name in classdict._member_names: del classdict[name] + # adjust the sunders + _order_ = classdict.pop('_order_', None) + # check for illegal enum names (any others?) - invalid_names = set(members) & {'mro', } + invalid_names = set(enum_members) & {'mro', } if invalid_names: raise ValueError('Invalid enum member name: {0}'.format( ','.join(invalid_names))) @@ -145,7 +167,7 @@ class EnumMeta(type): # a custom __new__ is doing something funky with the values -- such as # auto-numbering ;) for member_name in classdict._member_names: - value = members[member_name] + value = enum_members[member_name] if not isinstance(value, tuple): args = (value, ) else: @@ -159,7 +181,10 @@ class EnumMeta(type): else: enum_member = __new__(enum_class, *args) if not hasattr(enum_member, '_value_'): - enum_member._value_ = member_type(*args) + if member_type is object: + enum_member._value_ = value + else: + enum_member._value_ = member_type(*args) value = enum_member._value_ enum_member._name_ = member_name enum_member.__objclass__ = enum_class @@ -204,6 +229,14 @@ class EnumMeta(type): if save_new: enum_class.__new_member__ = __new__ enum_class.__new__ = Enum.__new__ + + # py3 support for definition order (helps keep py2/py3 code in sync) + if _order_ is not None: + if isinstance(_order_, str): + _order_ = _order_.replace(',', ' ').split() + if _order_ != enum_class._member_names_: + raise TypeError('member order does not match _order_') + return enum_class def __bool__(self): @@ -325,13 +358,18 @@ class EnumMeta(type): """ metacls = cls.__class__ bases = (cls, ) if type is None else (type, cls) + _, first_enum = cls._get_mixins_(bases) classdict = metacls.__prepare__(class_name, bases) # special processing needed for names? if isinstance(names, str): names = names.replace(',', ' ').split() if isinstance(names, (tuple, list)) and isinstance(names[0], str): - names = [(e, i) for (i, e) in enumerate(names, start)] + original_names, names = names, [] + last_value = None + for count, name in enumerate(original_names): + last_value = first_enum._generate_next_value_(name, start, count, last_value) + names.append((name, last_value)) # Here, names is either an iterable of (name, value) or a mapping. for item in names: @@ -473,6 +511,16 @@ class Enum(metaclass=EnumMeta): for member in cls._member_map_.values(): if member._value_ == value: return member + # still not found -- try _missing_ hook + return cls._missing_(value) + + @staticmethod + def _generate_next_value_(name, start, count, last_value): + if not count: + return start + return last_value + 1 + @classmethod + def _missing_(cls, value): raise ValueError("%r is not a valid %s" % (value, cls.__name__)) def __repr__(self): @@ -544,8 +592,14 @@ class Enum(metaclass=EnumMeta): source = vars(source) else: source = module_globals - members = {name: value for name, value in source.items() - if filter(name)} + # We use an OrderedDict of sorted source keys so that the + # _value2member_map is populated in the same order every time + # for a consistent reverse mapping of number to name when there + # are multiple names for the same number rather than varying + # between runs due to hash randomization of the module dictionary. + members = OrderedDict((name, source[name]) + for name in sorted(source.keys()) + if filter(name)) cls = cls(name, members, module=module) cls.__reduce_ex__ = _reduce_ex_by_name module_globals.update(cls.__members__) @@ -560,6 +614,184 @@ class IntEnum(int, Enum): def _reduce_ex_by_name(self, proto): return self.name +class Flags(Enum): + """Support for flags""" + @staticmethod + def _generate_next_value_(name, start, count, last_value): + """ + Generate the next value when not given. + + name: the name of the member + start: the initital start value or None + count: the number of existing members + last_value: the last value assigned or None + """ + if not count: + return start if start is not None else 1 + high_bit = _high_bit(last_value) + return 2 ** (high_bit+1) + + @classmethod + def _missing_(cls, value): + original_value = value + if value < 0: + value = ~value + possible_member = cls._create_pseudo_member_(value) + for member in possible_member._decompose_(): + if member._name_ is None and member._value_ != 0: + raise ValueError('%r is not a valid %s' % (original_value, cls.__name__)) + if original_value < 0: + possible_member = ~possible_member + return possible_member + + @classmethod + def _create_pseudo_member_(cls, value): + pseudo_member = cls._value2member_map_.get(value, None) + if pseudo_member is None: + # construct a non-singleton enum pseudo-member + pseudo_member = object.__new__(cls) + pseudo_member._name_ = None + pseudo_member._value_ = value + cls._value2member_map_[value] = pseudo_member + return pseudo_member + + def _decompose_(self): + """Extract all members from the value.""" + value = self._value_ + members = [] + cls = self.__class__ + for member in sorted(cls, key=lambda m: m._value_, reverse=True): + while _high_bit(value) > _high_bit(member._value_): + unknown = self._create_pseudo_member_(2 ** _high_bit(value)) + members.append(unknown) + value &= ~unknown._value_ + if ( + (value & member._value_ == member._value_) + and (member._value_ or not members) + ): + value &= ~member._value_ + members.append(member) + if not members or value: + members.append(self._create_pseudo_member_(value)) + members = list(members) + return members + + def __contains__(self, other): + if not isinstance(other, self.__class__): + return NotImplemented + return other._value_ & self._value_ == other._value_ + + def __iter__(self): + if self.value == 0: + return iter([]) + else: + return iter(self._decompose_()) + + def __repr__(self): + cls = self.__class__ + if self._name_ is not None: + return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_) + members = self._decompose_() + if len(members) == 1 and members[0]._name_ is None: + return '<%s: %r>' % (cls.__name__, members[0]._value_) + else: + return '<%s.%s: %r>' % ( + cls.__name__, + '|'.join([str(m._name_ or m._value_) for m in members]), + self._value_, + ) + + def __str__(self): + cls = self.__class__ + if self._name_ is not None: + return '%s.%s' % (cls.__name__, self._name_) + members = self._decompose_() + if len(members) == 1 and members[0]._name_ is None: + return '%s.%r' % (cls.__name__, members[0]._value_) + else: + return '%s.%s' % ( + cls.__name__, + '|'.join([str(m._name_ or m._value_) for m in members]), + ) + + def __or__(self, other): + if not isinstance(other, self.__class__): + return NotImplemented + return self.__class__(self._value_ | other._value_) + + def __and__(self, other): + if not isinstance(other, self.__class__): + return NotImplemented + return self.__class__(self._value_ & other._value_) + + def __xor__(self, other): + if not isinstance(other, self.__class__): + return NotImplemented + return self.__class__(self._value_ ^ other._value_) + + def __invert__(self): + members = self._decompose_() + inverted_members = [m for m in self.__class__ if m not in members and not m._value_ & self._value_] + inverted = reduce(_or_, inverted_members, self.__class__(0)) + return self.__class__(inverted) + + +class IntFlags(int, Flags): + """Support for integer-based Flags""" + + @classmethod + def _create_pseudo_member_(cls, value): + pseudo_member = cls._value2member_map_.get(value, None) + if pseudo_member is None: + # construct a non-singleton enum pseudo-member + pseudo_member = int.__new__(cls, value) + pseudo_member._name_ = None + pseudo_member._value_ = value + cls._value2member_map_[value] = pseudo_member + return pseudo_member + + @classmethod + def _missing_(cls, value): + possible_member = cls._create_pseudo_member_(value) + return possible_member + + def __or__(self, other): + if not isinstance(other, (self.__class__, int)): + return NotImplemented + return self.__class__(self._value_ | self.__class__(other)._value_) + + def __and__(self, other): + if not isinstance(other, (self.__class__, int)): + return NotImplemented + return self.__class__(self._value_ & self.__class__(other)._value_) + + def __xor__(self, other): + if not isinstance(other, (self.__class__, int)): + return NotImplemented + return self.__class__(self._value_ ^ self.__class__(other)._value_) + + __ror__ = __or__ + __rand__ = __and__ + __rxor__ = __xor__ + + def __invert__(self): + # members = self._decompose_() + # inverted_members = [m for m in self.__class__ if m not in members and not m._value_ & self._value_] + # inverted = reduce(_or_, inverted_members, self.__class__(0)) + return self.__class__(~self._value_) + + + + +def _high_bit(value): + """return the highest bit set in value""" + bit = 0 + while 'looking for the highest bit': + limit = 2 ** bit + if limit > value: + return bit - 1 + bit += 1 + def unique(enumeration): """Class decorator for enumerations ensuring unique member values.""" duplicates = [] |