Introducing __reduce_ex__, which is called with a protocol number argument

if it exists in preference over __reduce__.  Now Tim can go implement this
in cPickle.c.

3 files changed, 26 insertions, 15 deletions

diff --git a/Lib/copy_reg.py b/Lib/copy_reg.py
index 6dc52c2..fa4ce72 100644
--- a/Lib/copy_reg.py
+++ b/Lib/copy_reg.py
@@ -109,6 +109,17 @@ def _better_reduce(obj):
         dictitems = obj.iteritems()
     return __newobj__, (cls,) + args, state, listitems, dictitems
 
+# Extended reduce:
+
+def _reduce_ex(obj, proto=0):
+    obj_reduce = getattr(obj, "__reduce__", None)
+    if obj_reduce and obj.__class__.__reduce__ is not object.__reduce__:
+        return obj_reduce()
+    elif proto < 2:
+        return _reduce(obj)
+    else:
+        return _better_reduce(obj)
+
 def _slotnames(cls):
     """Return a list of slot names for a given class.
 
diff --git a/Lib/pickle.py b/Lib/pickle.py
index c62bddc..f997f38 100644
--- a/Lib/pickle.py
+++ b/Lib/pickle.py
@@ -304,21 +304,20 @@ class Pickler:
 
         # Check copy_reg.dispatch_table
         reduce = dispatch_table.get(t)
-        if not reduce:
-            # Check for a __reduce__ method.
-            # Subtle: get the unbound method from the class, so that
-            # protocol 2 can override the default __reduce__ that all
-            # classes inherit from object.  This has the added
-            # advantage that the call always has the form reduce(obj)
-            reduce = getattr(t, "__reduce__", None)
-            if self.proto >= 2:
-                # Protocol 2 can do better than the default __reduce__
-                if reduce is object.__reduce__:
-                    reduce = _better_reduce
-            if not reduce:
-                raise PicklingError("Can't pickle %r object: %r" %
-                                    (t.__name__, obj))
-        rv = reduce(obj)
+        if reduce:
+            rv = reduce(obj)
+        else:
+            # Check for a __reduce_ex__ method, fall back to __reduce__
+            reduce = getattr(obj, "__reduce_ex__", None)
+            if reduce:
+                rv = reduce(self.proto)
+            else:
+                reduce = getattr(obj, "__reduce__", None)
+                if reduce:
+                    rv = reduce()
+                else:
+                    raise PicklingError("Can't pickle %r object: %r" %
+                                        (t.__name__, obj))
 
         # Check for string returned by reduce(), meaning "save as global"
         if type(rv) is StringType:
diff --git a/Lib/test/test_descrtut.py b/Lib/test/test_descrtut.py
index 66cb749..5fa0897 100644
--- a/Lib/test/test_descrtut.py
+++ b/Lib/test/test_descrtut.py
@@ -210,6 +210,7 @@ Instead, you can get the same information from the list type:
      '__ne__',
      '__new__',
      '__reduce__',
+     '__reduce_ex__',
      '__repr__',
      '__rmul__',
      '__setattr__',