diff -r aab7258a31d3 Doc/library/enum.rst --- a/Doc/library/enum.rst Fri Feb 07 13:04:18 2014 -0500 +++ b/Doc/library/enum.rst Sat Feb 08 04:59:13 2014 -0800 @@ -362,24 +362,24 @@ Enumerations can be pickled and unpickle >>> from test.test_enum import Fruit >>> from pickle import dumps, loads >>> Fruit.tomato is loads(dumps(Fruit.tomato)) True The usual restrictions for pickling apply: picklable enums must be defined in the top level of a module, since unpickling requires them to be importable from that module. -.. warning:: +.. note:: - In order to support the singleton nature of enumeration members, pickle - protocol version 2 or higher must be used. + With pickle protocol version 4 it is possible to easily pickle enums + nested in other classes. Functional API -------------- The :class:`Enum` class is callable, providing the following functional API:: >>> Animal = Enum('Animal', 'ant bee cat dog') >>> Animal diff -r aab7258a31d3 Lib/enum.py --- a/Lib/enum.py Fri Feb 07 13:04:18 2014 -0500 +++ b/Lib/enum.py Sat Feb 08 04:59:13 2014 -0800 @@ -110,24 +110,28 @@ class EnumMeta(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 # Reverse value->name map for hashable values. enum_class._value2member_map_ = {} # check for a __getnewargs__, and if not present sabotage # pickling, since it won't work anyway - if (member_type is not object and - member_type.__dict__.get('__getnewargs__') is None - ): - _make_class_unpicklable(enum_class) + if member_type is not object: + for method in ( + '__getnewargs_ex__', '__getnewargs__', + '__reduce_ex__', '__reduce__'): + if member_type.__dict__.get(method) is not None: + break + else: + _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] if not isinstance(value, tuple): args = (value, ) else: @@ -159,57 +163,57 @@ class EnumMeta(type): try: # This may fail if value is not hashable. We can't add the value # to the map, and by-value lookups for this value will be # linear. enum_class._value2member_map_[value] = enum_member except TypeError: pass # double check that repr and friends are not the mixin's or various # things break (such as pickle) - for name in ('__repr__', '__str__', '__format__', '__getnewargs__'): + for name in ('__repr__', '__str__', '__format__', '__getnewargs__', '__reduce_ex__'): class_method = getattr(enum_class, name) obj_method = getattr(member_type, name, None) enum_method = getattr(first_enum, name, None) if obj_method is not None and obj_method is class_method: setattr(enum_class, name, enum_method) # replace any other __new__ with our own (as long as Enum is not None, # anyway) -- again, this is to support pickle if Enum is not None: # if the user defined their own __new__, save it before it gets # clobbered in case they subclass later if save_new: enum_class.__new_member__ = __new__ enum_class.__new__ = Enum.__new__ return enum_class - def __call__(cls, value, names=None, *, module=None, type=None): + def __call__(cls, value, names=None, *, module=None, qualname=None, type=None): """Either returns an existing member, or creates a new enum class. This method is used both when an enum class is given a value to match to an enumeration member (i.e. Color(3)) and for the functional API (i.e. Color = Enum('Color', names='red green blue')). When used for the functional API: `module`, if set, will be stored in the new class' __module__ attribute; `type`, if set, will be mixed in as the first base class. Note: if `module` is not set this routine will attempt to discover the calling module by walking the frame stack; if this is unsuccessful the resulting class will not be pickleable. """ if names is None: # simple value lookup return cls.__new__(cls, value) # otherwise, functional API: we're creating a new Enum type - return cls._create_(value, names, module=module, type=type) + return cls._create_(value, names, module=module, qualname=qualname, type=type) def __contains__(cls, member): return isinstance(member, cls) and member.name in cls._member_map_ def __delattr__(cls, attr): # nicer error message when someone tries to delete an attribute # (see issue19025). if attr in cls._member_map_: raise AttributeError( "%s: cannot delete Enum member." % cls.__name__) @@ -266,21 +270,21 @@ class EnumMeta(type): A simple assignment to the class namespace only changes one of the several possible ways to get an Enum member from the Enum class, resulting in an inconsistent Enumeration. """ member_map = cls.__dict__.get('_member_map_', {}) if name in member_map: raise AttributeError('Cannot reassign members.') super().__setattr__(name, value) - def _create_(cls, class_name, names=None, *, module=None, type=None): + def _create_(cls, class_name, names=None, *, module=None, qualname=None, type=None): """Convenience method to create a new Enum class. `names` can be: * A string containing member names, separated either with spaces or commas. Values are auto-numbered from 1. * An iterable of member names. Values are auto-numbered from 1. * An iterable of (member name, value) pairs. * A mapping of member name -> value. @@ -308,20 +312,22 @@ class EnumMeta(type): # module is ever developed if module is None: try: module = sys._getframe(2).f_globals['__name__'] except (AttributeError, ValueError) as exc: pass if module is None: _make_class_unpicklable(enum_class) else: enum_class.__module__ = module + if qualname is not None: + enum_class.__qualname__ = qualname return enum_class @staticmethod def _get_mixins_(bases): """Returns the type for creating enum members, and the first inherited enum class. bases: the tuple of bases that was given to __new__ @@ -461,20 +467,23 @@ class Enum(metaclass=EnumMeta): cls = self._member_type_ val = self.value return cls.__format__(val, format_spec) def __getnewargs__(self): return (self._value_, ) def __hash__(self): return hash(self._name_) + def __reduce_ex__(self, proto): + return self.__class__, self.__getnewargs__() + # DynamicClassAttribute is used to provide access to the `name` and # `value` properties of enum members while keeping some measure of # protection from modification, while still allowing for an enumeration # to have members named `name` and `value`. This works because enumeration # members are not set directly on the enum class -- __getattr__ is # used to look them up. @DynamicClassAttribute def name(self): """The name of the Enum member.""" diff -r aab7258a31d3 Lib/test/test_enum.py --- a/Lib/test/test_enum.py Fri Feb 07 13:04:18 2014 -0500 +++ b/Lib/test/test_enum.py Sat Feb 08 04:59:13 2014 -0800 @@ -45,37 +45,46 @@ except Exception as exc: try: Question = Enum('Question', 'who what when where why', module=__name__) except Exception as exc: Question = exc try: Answer = Enum('Answer', 'him this then there because') except Exception as exc: Answer = exc +try: + Theory = Enum('Theory', 'rule law supposition', qualname='spanish_inquisition') +except Exception as exc: + Theory = exc + # for doctests try: class Fruit(Enum): tomato = 1 banana = 2 cherry = 3 except Exception: pass -def test_pickle_dump_load(assertion, source, target=None): +def test_pickle_dump_load(assertion, source, target=None, + *, protocol=(0, HIGHEST_PROTOCOL)): + start, stop = protocol if target is None: target = source - for protocol in range(2, HIGHEST_PROTOCOL+1): + for protocol in range(start, stop+1): assertion(loads(dumps(source, protocol=protocol)), target) -def test_pickle_exception(assertion, exception, obj): - for protocol in range(2, HIGHEST_PROTOCOL+1): +def test_pickle_exception(assertion, exception, obj, + *, protocol=(0, HIGHEST_PROTOCOL)): + start, stop = protocol + for protocol in range(start, stop+1): with assertion(exception): dumps(obj, protocol=protocol) class TestHelpers(unittest.TestCase): # _is_descriptor, _is_sunder, _is_dunder def test_is_descriptor(self): class foo: pass for attr in ('__get__','__set__','__delete__'): @@ -533,26 +542,46 @@ class TestEnum(unittest.TestCase): raise Answer test_pickle_dump_load(self.assertIs, Answer.him) test_pickle_dump_load(self.assertIs, Answer) def test_pickle_enum_function_with_module(self): if isinstance(Question, Exception): raise Question test_pickle_dump_load(self.assertIs, Question.who) test_pickle_dump_load(self.assertIs, Question) + def test_enum_function_with_qualname(self): + if isinstance(Theory, Exception): + raise Theory + self.assertEqual(Theory.__qualname__, 'spanish_inquisition') + + def test_class_nested_enum_and_pickle_protocol_four(self): + # would normally just have this directly in the class namespace + class NestedEnum(Enum): + twigs = 'common' + shiny = 'rare' + + self.__class__.NestedEnum = NestedEnum + self.NestedEnum.__qualname__ = '%s.NestedEnum' % self.__class__.__name__ + test_pickle_exception( + self.assertRaises, PicklingError, self.NestedEnum.twigs, + protocol=(0, 3)) + test_pickle_dump_load(self.assertIs, self.NestedEnum.twigs, + protocol=(4, HIGHEST_PROTOCOL)) + def test_exploding_pickle(self): - BadPickle = Enum('BadPickle', 'dill sweet bread-n-butter') - BadPickle.__qualname__ = 'BadPickle' # needed for pickle protocol 4 + BadPickle = Enum( + 'BadPickle', 'dill sweet bread-n-butter', module=__name__) globals()['BadPickle'] = BadPickle - enum._make_class_unpicklable(BadPickle) # will overwrite __qualname__ - test_pickle_exception(self.assertRaises, TypeError, BadPickle.dill) + # now break BadPickle to test exception raising + enum._make_class_unpicklable(BadPickle) + test_pickle_exception(self.assertRaises, PicklingError, BadPickle.dill) test_pickle_exception(self.assertRaises, PicklingError, BadPickle) def test_string_enum(self): class SkillLevel(str, Enum): master = 'what is the sound of one hand clapping?' journeyman = 'why did the chicken cross the road?' apprentice = 'knock, knock!' self.assertEqual(SkillLevel.apprentice, 'knock, knock!') def test_getattr_getitem(self): @@ -920,20 +949,188 @@ class TestEnum(unittest.TestCase): self.assertIs(NEI.__new__, Enum.__new__) self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)") globals()['NamedInt'] = NamedInt globals()['NEI'] = NEI NI5 = NamedInt('test', 5) self.assertEqual(NI5, 5) test_pickle_dump_load(self.assertEqual, NI5, 5) self.assertEqual(NEI.y.value, 2) test_pickle_dump_load(self.assertIs, NEI.y) + def test_subclasses_with_getnewargs_ex(self): + class NamedInt(int): + __qualname__ = 'NamedInt' # needed for pickle protocol 4 + def __new__(cls, *args): + _args = args + name, *args = args + if len(args) == 0: + raise TypeError("name and value must be specified") + self = int.__new__(cls, *args) + self._intname = name + self._args = _args + return self + def __getnewargs_ex__(self): + return self._args, {} + @property + def __name__(self): + return self._intname + def __repr__(self): + # repr() is updated to include the name and type info + return "{}({!r}, {})".format(type(self).__name__, + self.__name__, + int.__repr__(self)) + def __str__(self): + # str() is unchanged, even if it relies on the repr() fallback + base = int + base_str = base.__str__ + if base_str.__objclass__ is object: + return base.__repr__(self) + return base_str(self) + # for simplicity, we only define one operator that + # propagates expressions + def __add__(self, other): + temp = int(self) + int( other) + if isinstance(self, NamedInt) and isinstance(other, NamedInt): + return NamedInt( + '({0} + {1})'.format(self.__name__, other.__name__), + temp ) + else: + return temp + + class NEI(NamedInt, Enum): + __qualname__ = 'NEI' # needed for pickle protocol 4 + x = ('the-x', 1) + y = ('the-y', 2) + + + self.assertIs(NEI.__new__, Enum.__new__) + self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)") + globals()['NamedInt'] = NamedInt + globals()['NEI'] = NEI + NI5 = NamedInt('test', 5) + self.assertEqual(NI5, 5) + test_pickle_dump_load(self.assertEqual, NI5, 5, protocol=(4, 4)) + self.assertEqual(NEI.y.value, 2) + test_pickle_dump_load(self.assertIs, NEI.y, protocol=(4, 4)) + + def test_subclasses_with_reduce(self): + class NamedInt(int): + __qualname__ = 'NamedInt' # needed for pickle protocol 4 + def __new__(cls, *args): + _args = args + name, *args = args + if len(args) == 0: + raise TypeError("name and value must be specified") + self = int.__new__(cls, *args) + self._intname = name + self._args = _args + return self + def __reduce__(self): + return self.__class__, self._args + @property + def __name__(self): + return self._intname + def __repr__(self): + # repr() is updated to include the name and type info + return "{}({!r}, {})".format(type(self).__name__, + self.__name__, + int.__repr__(self)) + def __str__(self): + # str() is unchanged, even if it relies on the repr() fallback + base = int + base_str = base.__str__ + if base_str.__objclass__ is object: + return base.__repr__(self) + return base_str(self) + # for simplicity, we only define one operator that + # propagates expressions + def __add__(self, other): + temp = int(self) + int( other) + if isinstance(self, NamedInt) and isinstance(other, NamedInt): + return NamedInt( + '({0} + {1})'.format(self.__name__, other.__name__), + temp ) + else: + return temp + + class NEI(NamedInt, Enum): + __qualname__ = 'NEI' # needed for pickle protocol 4 + x = ('the-x', 1) + y = ('the-y', 2) + + + self.assertIs(NEI.__new__, Enum.__new__) + self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)") + globals()['NamedInt'] = NamedInt + globals()['NEI'] = NEI + NI5 = NamedInt('test', 5) + self.assertEqual(NI5, 5) + test_pickle_dump_load(self.assertEqual, NI5, 5) + self.assertEqual(NEI.y.value, 2) + test_pickle_dump_load(self.assertIs, NEI.y) + + def test_subclasses_with_reduce_ex(self): + class NamedInt(int): + __qualname__ = 'NamedInt' # needed for pickle protocol 4 + def __new__(cls, *args): + _args = args + name, *args = args + if len(args) == 0: + raise TypeError("name and value must be specified") + self = int.__new__(cls, *args) + self._intname = name + self._args = _args + return self + def __reduce_ex__(self, proto): + return self.__class__, self._args + @property + def __name__(self): + return self._intname + def __repr__(self): + # repr() is updated to include the name and type info + return "{}({!r}, {})".format(type(self).__name__, + self.__name__, + int.__repr__(self)) + def __str__(self): + # str() is unchanged, even if it relies on the repr() fallback + base = int + base_str = base.__str__ + if base_str.__objclass__ is object: + return base.__repr__(self) + return base_str(self) + # for simplicity, we only define one operator that + # propagates expressions + def __add__(self, other): + temp = int(self) + int( other) + if isinstance(self, NamedInt) and isinstance(other, NamedInt): + return NamedInt( + '({0} + {1})'.format(self.__name__, other.__name__), + temp ) + else: + return temp + + class NEI(NamedInt, Enum): + __qualname__ = 'NEI' # needed for pickle protocol 4 + x = ('the-x', 1) + y = ('the-y', 2) + + + self.assertIs(NEI.__new__, Enum.__new__) + self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)") + globals()['NamedInt'] = NamedInt + globals()['NEI'] = NEI + NI5 = NamedInt('test', 5) + self.assertEqual(NI5, 5) + test_pickle_dump_load(self.assertEqual, NI5, 5) + self.assertEqual(NEI.y.value, 2) + test_pickle_dump_load(self.assertIs, NEI.y) + def test_subclasses_without_getnewargs(self): class NamedInt(int): __qualname__ = 'NamedInt' def __new__(cls, *args): _args = args name, *args = args if len(args) == 0: raise TypeError("name and value must be specified") self = int.__new__(cls, *args) self._intname = name @@ -970,21 +1167,21 @@ class TestEnum(unittest.TestCase): x = ('the-x', 1) y = ('the-y', 2) self.assertIs(NEI.__new__, Enum.__new__) self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)") globals()['NamedInt'] = NamedInt globals()['NEI'] = NEI NI5 = NamedInt('test', 5) self.assertEqual(NI5, 5) self.assertEqual(NEI.y.value, 2) - test_pickle_exception(self.assertRaises, TypeError, NEI.x) + test_pickle_exception(self.assertRaises, PicklingError, NEI.x) test_pickle_exception(self.assertRaises, PicklingError, NEI) def test_tuple_subclass(self): class SomeTuple(tuple, Enum): __qualname__ = 'SomeTuple' # needed for pickle protocol 4 first = (1, 'for the money') second = (2, 'for the show') third = (3, 'for the music') self.assertIs(type(SomeTuple.first), SomeTuple) self.assertIsInstance(SomeTuple.second, tuple)