Close issue20653: allow Enum subclasses to override __reduce_ex__

ethanfurman · ethanfurman · commit 9a0cbcc4f8f6 · 2014-02-18T12:37:12.000-08:00
diff --git a/Lib/enum.py b/Lib/enum.py
@@ -116,12 +116,14 @@ def __new__(metacls, cls, bases, classdict):
         enum_class._value2member_map_ = {}
 
         # check for a supported pickle protocols, and if not present sabotage
-        # pickling, since it won't work anyway
-        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)
+        # pickling, since it won't work anyway.
+        # if new class implements its own __reduce_ex__, do not sabotage
+        if classdict.get('__reduce_ex__') is None:
+            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
@@ -167,7 +169,7 @@ def __new__(metacls, cls, bases, classdict):
 
         # 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__', '__reduce_ex__'):
+        for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
             class_method = getattr(enum_class, name)
             obj_method = getattr(member_type, name, None)
             enum_method = getattr(first_enum, name, None)
@@ -192,8 +194,9 @@ def __call__(cls, value, names=None, *, module=None, qualname=None, type=None):
         (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.
+        the new class' __module__ attribute; `qualname`, if set, will be stored
+        in the new class' __qualname__ 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
@@ -465,14 +468,11 @@ def __format__(self, format_spec):
             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__()
+        return self.__class__, (self._value_, )
 
     # DynamicClassAttribute is used to provide access to the `name` and
     # `value` properties of enum members while keeping some measure of
diff --git a/Lib/test/test_enum.py b/Lib/test/test_enum.py
@@ -956,6 +956,7 @@ class NEI(NamedInt, Enum):
         test_pickle_dump_load(self.assertEqual, NI5, 5)
         self.assertEqual(NEI.y.value, 2)
         test_pickle_dump_load(self.assertIs, NEI.y)
+        test_pickle_dump_load(self.assertIs, NEI)
 
     def test_subclasses_with_getnewargs_ex(self):
         class NamedInt(int):
@@ -1012,6 +1013,7 @@ class NEI(NamedInt, Enum):
         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))
+        test_pickle_dump_load(self.assertIs, NEI)
 
     def test_subclasses_with_reduce(self):
         class NamedInt(int):
@@ -1068,6 +1070,7 @@ class NEI(NamedInt, Enum):
         test_pickle_dump_load(self.assertEqual, NI5, 5)
         self.assertEqual(NEI.y.value, 2)
         test_pickle_dump_load(self.assertIs, NEI.y)
+        test_pickle_dump_load(self.assertIs, NEI)
 
     def test_subclasses_with_reduce_ex(self):
         class NamedInt(int):
@@ -1124,8 +1127,9 @@ class NEI(NamedInt, Enum):
         test_pickle_dump_load(self.assertEqual, NI5, 5)
         self.assertEqual(NEI.y.value, 2)
         test_pickle_dump_load(self.assertIs, NEI.y)
+        test_pickle_dump_load(self.assertIs, NEI)
 
-    def test_subclasses_without_getnewargs(self):
+    def test_subclasses_without_direct_pickle_support(self):
         class NamedInt(int):
             __qualname__ = 'NamedInt'
             def __new__(cls, *args):
@@ -1178,6 +1182,61 @@ class NEI(NamedInt, Enum):
         test_pickle_exception(self.assertRaises, TypeError, NEI.x)
         test_pickle_exception(self.assertRaises, PicklingError, NEI)
 
+    def test_subclasses_without_direct_pickle_support_using_name(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
+                self._args = _args
+                return self
+            @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'
+            x = ('the-x', 1)
+            y = ('the-y', 2)
+            def __reduce_ex__(self, proto):
+                return getattr, (self.__class__, self._name_)
+
+        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_dump_load(self.assertIs, NEI.y)
+        test_pickle_dump_load(self.assertIs, NEI)
+
     def test_tuple_subclass(self):
         class SomeTuple(tuple, Enum):
             __qualname__ = 'SomeTuple'      # needed for pickle protocol 4
