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.

4 files changed, 36 insertions, 19 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__',
diff --git a/Objects/typeobject.c b/Objects/typeobject.c
index e238056..a1e5b88 100644
--- a/Objects/typeobject.c
+++ b/Objects/typeobject.c
@@ -2441,11 +2441,15 @@ static PyGetSetDef object_getsets[] = {
 };
 
 static PyObject *
-object_reduce(PyObject *self, PyObject *args)
+object_reduce_ex(PyObject *self, PyObject *args)
 {
-	/* Call copy_reg._reduce(self) */
+	/* Call copy_reg._reduce_ex(self, proto) */
 	static PyObject *copy_reg_str;
 	PyObject *copy_reg, *res;
+	int proto = 0;
+
+	if (!PyArg_ParseTuple(args, "|i:__reduce_ex__", &proto))
+		return NULL;
 
 	if (!copy_reg_str) {
 		copy_reg_str = PyString_InternFromString("copy_reg");
@@ -2455,13 +2459,15 @@ object_reduce(PyObject *self, PyObject *args)
 	copy_reg = PyImport_Import(copy_reg_str);
 	if (!copy_reg)
 		return NULL;
-	res = PyEval_CallMethod(copy_reg, "_reduce", "(O)", self);
+	res = PyEval_CallMethod(copy_reg, "_reduce_ex", "(Oi)", self, proto);
 	Py_DECREF(copy_reg);
 	return res;
 }
 
 static PyMethodDef object_methods[] = {
-	{"__reduce__", object_reduce, METH_NOARGS,
+	{"__reduce_ex__", object_reduce_ex, METH_VARARGS,
+	 PyDoc_STR("helper for pickle")},
+	{"__reduce__", object_reduce_ex, METH_VARARGS,
 	 PyDoc_STR("helper for pickle")},
 	{0}
 };