diff --git a/Lib/enum.py b/Lib/enum.py --- a/Lib/enum.py +++ b/Lib/enum.py @@ -1,21 +1,24 @@ import sys from types import MappingProxyType, DynamicClassAttribute # try _collections first to reduce startup cost try: - from _collections import OrderedDict + from _collections import OrderedDict, defaultdict except ImportError: - from collections import OrderedDict + from collections import OrderedDict, defaultdictt -__all__ = ['EnumMeta', 'Enum', 'IntEnum', 'unique'] +__all__ = [ + 'EnumMeta', 'Enum', 'IntEnum', 'AutoEnum', 'AutoNameEnum', 'UniqueEnum', + 'unique', 'AutoNumber', 'AutoName', 'Unique', + ] def _is_descriptor(obj): """Returns True if obj is a descriptor, False otherwise.""" return ( hasattr(obj, '__get__') or hasattr(obj, '__set__') or hasattr(obj, '__delete__')) @@ -44,90 +47,227 @@ def _make_class_unpicklable(cls): class _EnumDict(dict): """Track enum member order and ensure member names are not reused. EnumMeta will use the names found in self._member_names as the enumeration member names. """ def __init__(self): - super().__init__() + super(_EnumDict, self).__init__() + # list of enum members self._member_names = [] + # starting value for AutoNumber + self._value = None + # when the magic turns off + self._locked = True + # list of temporary names + self._ignore = [] + # if _sunder_ values can be changed via the class body + self._init = True + # which auto, if any, is active + self._auto = None + # unique? + self._unique = False + + def __getitem__(self, key): + if ( + self._locked + or key in self + or key in self._ignore + or _is_sunder(key) + or _is_dunder(key) + ): + return super(_EnumDict, self).__getitem__(key) + if self._auto is AutoNumber: + try: + # try to generate the next value + value = self._value + 1 + self.__setitem__(key, value) + except: + # couldn't work the magic, report error + raise ValueError('Unable to generate value for %r from %r' % (key, self._value)) + elif self._auto is AutoName: + value = key.lower() + self.__setitem__(key, value) + else: + raise KeyError('%s not found' % key) + return value def __setitem__(self, key, value): """Changes anything not dundered or not a descriptor. If an enum member name is used twice, an error is raised; duplicate values are not checked for. Single underscore (sunder) names are reserved. """ if _is_sunder(key): - raise ValueError('_names_ are reserved for future Enum use') + if key not in ('_settings_', '_order_', '_ignore_', '_start_'): + raise ValueError('_names_ are reserved for future Enum use') + elif not self._init: + raise ValueError('cannot set %r after init phase' % key) + elif key == '_ignore_': + if isinstance(value, str): + value = value.split() + else: + value = list(value) + self._ignore = value + already = set(value) & set(self._member_names) + if already: + raise ValueError('_ignore_ cannot specify already set names: %r' % (already, )) + elif key == '_start_': + self._value = value - 1 + self._locked = False + elif key == '_settings_': + if not isinstance(value, tuple): + value = value, + if AutoName in value and AutoNumber in value: + raise TypeError('cannot specify both AutoName and AutoNumber') + allowed_settings = dict.fromkeys(['autonumber', 'autoname', 'unique']) + for arg in value: + if arg not in allowed_settings: + raise TypeError('unknown setting: %r' % arg) + allowed_settings[arg] = True + if arg in (AutoName, AutoNumber): + self._auto = arg + self._locked = allowed_settings['autonumber'] and allowed_settings['autoname'] elif _is_dunder(key): - pass + if key == '__order__': + key = '_order_' + if _is_descriptor(value): + self._locked = True elif key in self._member_names: # descriptor overwriting an enum? raise TypeError('Attempted to reuse key: %r' % key) + elif key in self._ignore: + pass elif not _is_descriptor(value): + self._init = False 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) + if isinstance(value, int): + self._value = value + else: + # not a new member, turn off the autoassign magic + self._locked = True + self._init = False super().__setitem__(key, value) # Dummy value for Enum as EnumMeta explicitly checks for it, but of course # until EnumMeta finishes running the first time the Enum class doesn't exist. # This is also why there are checks in EnumMeta like `if Enum is not None` Enum = None +Unique = 'unique' +AutoName = 'autoname' +AutoNumber = 'autonumber' - +_ignore_sentinel = object() class EnumMeta(type): """Metaclass for Enum""" @classmethod - def __prepare__(metacls, cls, bases): - return _EnumDict() + def __prepare__(metacls, cls, bases, start=None, settings=(), ignore=_ignore_sentinel): + if not isinstance(settings, tuple): + settings = (settings, ) + # inherit previous flags + member_type, first_enum = metacls._get_mixins_(bases) + if first_enum is not None: + first_settings = getattr(first_enum, '_settings_', ()) + if start is None: + start = getattr(first_enum, '_start_', None) + else: + first_settings = () + enum_dict = _EnumDict() + # combine found flags and settings + settings = tuple(set(settings+first_settings)) + if AutoNumber in settings and start is None: + start = 1 + if settings: + # and add to class dict + enum_dict['_settings_'] = settings + if ignore is _ignore_sentinel: + enum_dict['_ignore_'] = 'property classmethod staticmethod'.split() + elif ignore: + enum_dict['_ignore_'] = ignore + if start is not None: + enum_dict['_start_'] = start + return enum_dict - def __new__(metacls, cls, bases, classdict): + def __init__(cls, *args , **kwds): + super(EnumMeta, cls).__init__(*args) + + 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 # class will fail). member_type, first_enum = metacls._get_mixins_(bases) __new__, save_new, use_args = metacls._find_new_(classdict, member_type, first_enum) # 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 + _settings_ = classdict.get('_settings_', ()) + _order_ = classdict.pop('_order_', None) + classdict.pop('_ignore_', None) + # maybe check for uniqueness of values + if Unique in _settings_: + value_map = defaultdict(list) + for member, value in enum_members.items(): + value_map[value].append(member) + dups = [] + for value, members in value_map.items(): + if len(members) > 1: + dups.append('members %r all have the value of %r' % (members, value)) + if dups: + # use ValueError for consistency with unique decorator + raise ValueError('Uniqueness invalid for %s:\n %s' % (cls, '\n '.join(dups))) + + # 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() + unique_members = [n for n in clsdict._member_names if n in _order_] + if _order_ != unique_members: + raise TypeError('member order does not match _order_') + # 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))) # create a default docstring if one has not been provided if '__doc__' not in classdict: classdict['__doc__'] = 'An enumeration.' # 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 enum_class._member_type_ = member_type + _start_ = classdict._value + if _start_ is not None: + # should be number of next member + _start_ += 1 + enum_class._start_ = _start_ # save attributes from super classes so we know if we can take # the shortcut of storing members in the class dict base_attributes = {a for b in enum_class.mro() for a in b.__dict__} # Reverse value->name map for hashable values. enum_class._value2member_map_ = {} # If a custom type is mixed into the Enum, and it does not know how # to pickle itself, pickle.dumps will succeed but pickle.loads will @@ -144,21 +284,21 @@ class EnumMeta(type): methods = ('__getnewargs_ex__', '__getnewargs__', '__reduce_ex__', '__reduce__') if not any(m in member_type.__dict__ for m in methods): _make_class_unpicklable(enum_class) # instantiate them, checking for duplicates as we go # we instantiate first instead of checking for duplicates first in case # 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: args = value if member_type is tuple: # special case for tuple enums args = (args, ) # wrap it one more time if not use_args: enum_member = __new__(enum_class) if not hasattr(enum_member, '_value_'): enum_member._value_ = value @@ -565,20 +705,36 @@ class Enum(metaclass=EnumMeta): return cls class IntEnum(int, Enum): """Enum where members are also (and must be) ints""" def _reduce_ex_by_name(self, proto): return self.name +class _EnumConstants(str, Enum): + """EnumConstants for specifying Enum settings.""" + AutoNumber = 'autonumber' + AutoName = 'autoname' + Unique = 'unique' +globals().update(_EnumConstants.__members__) + +class AutoEnum(Enum, settings=AutoNumber): + """Enum where values are automatically assigned.""" + +class AutoNameEnum(Enum, settings=AutoName): + """Enum where member values are lower-cased versions of the name.""" + +class UniqueEnum(Enum, settings=Unique): + """Enum where all values must be unique (no aliases allowed).""" + def unique(enumeration): """Class decorator for enumerations ensuring unique member values.""" duplicates = [] for name, member in enumeration.__members__.items(): if name != member.name: duplicates.append((name, member.name)) if duplicates: alias_details = ', '.join( ["%s -> %s" % (alias, name) for (alias, name) in duplicates]) raise ValueError('duplicate values found in %r: %s' % diff --git a/Lib/test/test_enum.py b/Lib/test/test_enum.py --- a/Lib/test/test_enum.py +++ b/Lib/test/test_enum.py @@ -1,16 +1,16 @@ import enum import inspect import pydoc import unittest from collections import OrderedDict -from enum import Enum, IntEnum, EnumMeta, unique +from enum import Enum, IntEnum, AutoEnum, AutoNameEnum, UniqueEnum, EnumMeta, unique from io import StringIO from pickle import dumps, loads, PicklingError, HIGHEST_PROTOCOL from test import support # for pickle tests try: class Stooges(Enum): LARRY = 1 CURLY = 2 MOE = 3 @@ -1564,22 +1564,180 @@ class TestEnum(unittest.TestCase): class LabelledList(LabelledIntEnum): unprocessed = (1, "Unprocessed") payment_complete = (2, "Payment Complete") self.assertEqual(list(LabelledList), [LabelledList.unprocessed, LabelledList.payment_complete]) self.assertEqual(LabelledList.unprocessed, 1) self.assertEqual(LabelledList(1), LabelledList.unprocessed) +class TestAutoNumber(unittest.TestCase): + + def test_autonumbering(self): + class Color(AutoEnum): + red + green + blue + self.assertEqual(list(Color), [Color.red, Color.green, Color.blue]) + self.assertEqual(Color.red.value, 1) + self.assertEqual(Color.green.value, 2) + self.assertEqual(Color.blue.value, 3) + + def test_badly_overridden_ignore(self): + with self.assertRaisesRegex(TypeError, "'int' object is not callable"): + class Color(AutoEnum): + _ignore_ = () + red + green + blue + @property + def whatever(self): + pass + with self.assertRaisesRegex(TypeError, "'int' object is not callable"): + class Color(AutoEnum, ignore=None): + red + green + blue + @property + def whatever(self): + pass + + def test_property(self): + class Color(AutoEnum): + _ignore_ = 'property' + red + green + blue + @property + def cap_name(self): + return self.name.title() + self.assertEqual(Color.blue.cap_name, 'Blue') + + def test_magic_turns_off(self): + with self.assertRaises(NameError): + class Color(AutoEnum): + _ignore_ = 'property' + red + green + blue + @property + def cap_name(self): + return self.name.title() + brown + + with self.assertRaises(NameError): + class Color(Enum, settings=AutoNumber): + _ignore_ = 'property' + red + green + blue + def hello(self): + print('Hello! My serial is %s.' % self.value) + rose + + with self.assertRaises(NameError): + class Color(Enum, settings=AutoNumber): + _ignore_ = 'property' + red + green + blue + def __init__(self, *args): + pass + rose + + +class TestAutoName(unittest.TestCase): + + def test_autonaming(self): + class Color(AutoNameEnum): + Red + Green + Blue + self.assertEqual(list(Color), [Color.Red, Color.Green, Color.Blue]) + self.assertEqual(Color.Red.value, 'red') + self.assertEqual(Color.Green.value, 'green') + self.assertEqual(Color.Blue.value, 'blue') + + def test_overridden_ignore(self): + with self.assertRaisesRegex(TypeError, "'str' object is not callable"): + class Color(AutoNameEnum): + _ignore_ = () + red + green + blue + @property + def whatever(self): + pass + with self.assertRaisesRegex(TypeError, "'str' object is not callable"): + class Color(AutoNameEnum, ignore=None): + red + green + blue + @property + def whatever(self): + pass + + def test_property(self): + class Color(AutoNameEnum): + red + green + blue + @property + def upper_name(self): + return self.name.upper() + self.assertEqual(Color.blue.upper_name, 'BLUE') + + def test_magic_turns_off(self): + with self.assertRaises(NameError): + class Color(AutoNameEnum): + red + green + blue + @property + def cap_name(self): + return self.name.title() + brown + + with self.assertRaises(NameError): + class Color(Enum, settings=AutoNumber): + red + green + blue + def hello(self): + print('Hello! My value %s.' % self.value) + rose + + with self.assertRaises(NameError): + class Color(Enum, settings=AutoNumber): + red + green + blue + def __init__(self, *args): + pass + rose + + class TestUnique(unittest.TestCase): + def test_unique_in_settings(self): + class Nice(UniqueEnum): + red = 1 + green = 2 + blue = 3 + with self.assertRaisesRegex(ValueError, 'all have the value of 2'): + class NotNice(UniqueEnum): + red = 1 + green = 2 + blue = 3 + grene = 2 + def test_unique_clean(self): @unique class Clean(Enum): one = 1 two = 'dos' tres = 4.0 @unique class Cleaner(IntEnum): single = 1 double = 2