diff --git a/Doc/library/enum.rst b/Doc/library/enum.rst index 13f8a3c4e2e..7f800e36a0d 100644 --- a/Doc/library/enum.rst +++ b/Doc/library/enum.rst @@ -369,10 +369,10 @@ 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 @@ -420,6 +420,14 @@ The solution is to specify the module name explicitly as follows:: >>> Animals = Enum('Animals', 'ant bee cat dog', module=__name__) +The new pickle protocol 4 also, in some circumstances, relies on +:attr:``__qualname__`` being set to the location where pickle will be able +to find the class. For example, if the class was made available in class +SomeData in the global scope:: + + >>> Animals = Enum('Animals', 'ant bee cat dog', qualname='SomeData.Animals') + + Derived Enumerations -------------------- diff --git a/Lib/enum.py b/Lib/enum.py index 7d58f8d17a6..794f68e60ca 100644 --- a/Lib/enum.py +++ b/Lib/enum.py @@ -31,9 +31,9 @@ def _is_sunder(name): def _make_class_unpicklable(cls): """Make the given class un-picklable.""" - def _break_on_call_reduce(self): + def _break_on_call_reduce(self, proto): raise TypeError('%r cannot be pickled' % self) - cls.__reduce__ = _break_on_call_reduce + cls.__reduce_ex__ = _break_on_call_reduce cls.__module__ = '' @@ -115,12 +115,13 @@ class EnumMeta(type): # Reverse value->name map for hashable values. enum_class._value2member_map_ = {} - # check for a __getnewargs__, and if not present sabotage + # check for a supported pickle protocols, 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: + methods = ('__getnewargs_ex__', '__getnewargs__', + '__reduce_ex__', '__reduce__') + if not any(map(member_type.__dict__.get, 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 @@ -166,7 +167,7 @@ class EnumMeta(type): # 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) @@ -183,7 +184,7 @@ class EnumMeta(type): 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 @@ -202,7 +203,7 @@ class EnumMeta(type): 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_ @@ -273,7 +274,7 @@ class EnumMeta(type): 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: @@ -315,6 +316,8 @@ class EnumMeta(type): _make_class_unpicklable(enum_class) else: enum_class.__module__ = module + if qualname is not None: + enum_class.__qualname__ = qualname return enum_class @@ -468,6 +471,9 @@ class Enum(metaclass=EnumMeta): 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 diff --git a/Lib/test/test_enum.py b/Lib/test/test_enum.py index dfade178733..02009afbc85 100644 --- a/Lib/test/test_enum.py +++ b/Lib/test/test_enum.py @@ -52,6 +52,11 @@ try: 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): @@ -61,14 +66,18 @@ try: 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) @@ -101,6 +110,7 @@ class TestHelpers(unittest.TestCase): class TestEnum(unittest.TestCase): + def setUp(self): class Season(Enum): SPRING = 1 @@ -540,11 +550,31 @@ class TestEnum(unittest.TestCase): 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__ + # now break BadPickle to test exception raising + enum._make_class_unpicklable(BadPickle) test_pickle_exception(self.assertRaises, TypeError, BadPickle.dill) test_pickle_exception(self.assertRaises, PicklingError, BadPickle) @@ -927,6 +957,174 @@ class TestEnum(unittest.TestCase): 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'