update pickle protocol

Author: brijkapadia

Lib/functools.py

@@ -398,6 +398,12 @@ def __get__(self, obj, objtype=None):
             return self
         return MethodType(self, obj)
 
+    def __reduce_ex__(self, protocol):
+        if protocol >= 2:
+            return self.__reduce__()
+        return type(self), (self.func,), (self.func, self.args,
+               dict(self.keywords) or None, self.__dict__ or None)
+
     def __reduce__(self):
         return type(self), (self.func,), (self.func, self.args,
                self.keywords or None, self.__dict__ or None)

Lib/test/test_functools.py

@@ -62,7 +62,6 @@ class BadTuple(tuple):
     def __add__(self, other):
         return list(self) + list(other)
 
-
 class MyDict(frozendict):
     pass
 
@@ -342,7 +341,7 @@ def test_pickle(self):
         with replaced_module('functools', self.module):
             f = self.partial(signature, ['asdf'], bar=[True])
             f.attr = []
-            for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
+            for proto in range(pickle.HIGHEST_PROTOCOL + 1):
                 f_copy = pickle.loads(pickle.dumps(f, proto))
                 self.assertEqual(signature(f_copy), signature(f))
 
@@ -470,7 +469,7 @@ def test_recursive_pickle(self):
             f = self.partial(capture)
             f.__setstate__((f, (), {}, {}))
             try:
-                for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
+                for proto in range(pickle.HIGHEST_PROTOCOL + 1):
                     # gh-117008: Small limit since pickle uses C stack memory
                     with support.infinite_recursion(100):
                         with self.assertRaises(RecursionError):
@@ -481,7 +480,7 @@ def test_recursive_pickle(self):
             f = self.partial(capture)
             f.__setstate__((capture, (f,), {}, {}))
             try:
-                for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
+                for proto in range(pickle.HIGHEST_PROTOCOL + 1):
                     f_copy = pickle.loads(pickle.dumps(f, proto))
                     try:
                         self.assertIs(f_copy.args[0], f_copy)
@@ -493,7 +492,7 @@ def test_recursive_pickle(self):
             f = self.partial(capture)
             f.__setstate__((capture, (), {'a': f}, {}))
             try:
-                for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
+                for proto in range(pickle.HIGHEST_PROTOCOL + 1):
                     f_copy = pickle.loads(pickle.dumps(f, proto))
                     try:
                         self.assertIs(f_copy.keywords['a'], f_copy)

Modules/_functoolsmodule.c

@@ -777,6 +777,10 @@ partial_repr(PyObject *self)
    operation so we define a __setstate__ that replaces all the information
    about the partial.  If we only replaced part of it someone would use
    it as a hook to do strange things.
+
+   Additionally, since frozendict does not work for pickle protocols 0 and 1,
+   __reduce_ex__ creates a temporary dict for protocols 0 and 1 and calls
+   __reduce__ for protocols 2+.
  */
 
 static PyObject *
@@ -788,6 +792,28 @@ partial_reduce(PyObject *self, PyObject *Py_UNUSED(args))
                          pto->dict ? pto->dict : Py_None);
 }
 
+static PyObject *
+partial_reduce_ex(PyObject *self, PyObject *args)
+{
+    int64_t protocol;
+
+    if (!PyArg_ParseTuple(args, "l", &protocol)) {
+        return NULL;
+    }
+
+    if (protocol >= 2) {
+        return partial_reduce(self, NULL);
+    }
+
+    partialobject *pto = partialobject_CAST(self);
+    PyObject *keywords_dict = PyObject_CallOneArg((PyObject*)&PyDict_Type, pto->kw);
+    PyObject *result = Py_BuildValue("O(O)(OOOO)", Py_TYPE(pto), pto->fn, pto->fn,
+                         pto->args, keywords_dict,
+                         pto->dict ? pto->dict : Py_None);
+    Py_DECREF(keywords_dict);
+    return result;
+}
+
 static PyObject *
 partial_setstate(PyObject *self, PyObject *state)
 {
@@ -878,6 +904,7 @@ partial_setstate(PyObject *self, PyObject *state)
 
 static PyMethodDef partial_methods[] = {
     {"__reduce__", partial_reduce, METH_NOARGS},
+    {"__reduce_ex__", partial_reduce_ex, METH_VARARGS},
     {"__setstate__", partial_setstate, METH_O},
     {"__class_getitem__",    Py_GenericAlias,
     METH_O|METH_CLASS,       PyDoc_STR("See PEP 585")},