Close issue20534: all pickle protocols now supported.

1 files changed, 205 insertions, 7 deletions

diff --git a/Lib/test/test_enum.py b/Lib/test/test_enum.py
index dfade17..02009af 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'