diff --git a/Lib/test/test_functools.py b/Lib/test/test_functools.py index 30d419da21d..cd645d34bf2 100644 --- a/Lib/test/test_functools.py +++ b/Lib/test/test_functools.py @@ -2,6 +2,7 @@ import functools import unittest from test import test_support from weakref import proxy +import pickle @staticmethod def PythonPartial(func, *args, **keywords): @@ -19,6 +20,10 @@ def capture(*args, **kw): """capture all positional and keyword arguments""" return args, kw +def signature(part): + """ return the signature of a partial object """ + return (part.func, part.args, part.keywords, part.__dict__) + class TestPartial(unittest.TestCase): thetype = functools.partial @@ -140,6 +145,12 @@ class TestPartial(unittest.TestCase): join = self.thetype(''.join) self.assertEqual(join(data), '0123456789') + def test_pickle(self): + f = self.thetype(signature, 'asdf', bar=True) + f.add_something_to__dict__ = True + f_copy = pickle.loads(pickle.dumps(f)) + self.assertEqual(signature(f), signature(f_copy)) + class PartialSubclass(functools.partial): pass @@ -147,11 +158,13 @@ class TestPartialSubclass(TestPartial): thetype = PartialSubclass - class TestPythonPartial(TestPartial): thetype = PythonPartial + # the python version isn't picklable + def test_pickle(self): pass + class TestUpdateWrapper(unittest.TestCase): def check_wrapper(self, wrapper, wrapped, diff --git a/Misc/ACKS b/Misc/ACKS index d5b055cd902..b7cd25a75b1 100644 --- a/Misc/ACKS +++ b/Misc/ACKS @@ -168,6 +168,7 @@ Roger Dev Raghuram Devarakonda Toby Dickenson Mark Dickinson +Jack Diederich Yves Dionne Daniel Dittmar Jaromir Dolecek diff --git a/Misc/NEWS b/Misc/NEWS index 3d15073ef8e..62b37f211a5 100644 --- a/Misc/NEWS +++ b/Misc/NEWS @@ -710,6 +710,8 @@ Extension Modules - Issue #4396: The parser module now correctly validates the with statement. +- Issue #5228: Make functools.partial objects can now be pickled. + Tests ----- diff --git a/Modules/_functoolsmodule.c b/Modules/_functoolsmodule.c index 5dfbc15c139..92fa4cb955d 100644 --- a/Modules/_functoolsmodule.c +++ b/Modules/_functoolsmodule.c @@ -274,6 +274,53 @@ static PyGetSetDef partial_getsetlist[] = { {NULL} /* Sentinel */ }; +/* Pickle strategy: + __reduce__ by itself doesn't support getting kwargs in the unpickle + 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 stange things. + */ + +PyObject * +partial_reduce(partialobject *pto, PyObject *unused) +{ + return Py_BuildValue("O(O)(OOOO)", Py_TYPE(pto), pto->fn, pto->fn, + pto->args, pto->kw, + pto->dict ? pto->dict : Py_None); +} + +PyObject * +partial_setstate(partialobject *pto, PyObject *args) +{ + PyObject *fn, *fnargs, *kw, *dict; + if (!PyArg_ParseTuple(args, "(OOOO):__setstate__", + &fn, &fnargs, &kw, &dict)) + return NULL; + Py_XDECREF(pto->fn); + Py_XDECREF(pto->args); + Py_XDECREF(pto->kw); + Py_XDECREF(pto->dict); + pto->fn = fn; + pto->args = fnargs; + pto->kw = kw; + if (dict != Py_None) { + pto->dict = dict; + Py_INCREF(dict); + } else { + pto->dict = NULL; + } + Py_INCREF(fn); + Py_INCREF(fnargs); + Py_INCREF(kw); + Py_RETURN_NONE; +} + +static PyMethodDef partial_methods[] = { + {"__reduce__", (PyCFunction)partial_reduce, METH_NOARGS}, + {"__setstate__", (PyCFunction)partial_setstate, METH_VARARGS}, + {NULL, NULL} /* sentinel */ +}; + static PyTypeObject partial_type = { PyVarObject_HEAD_INIT(NULL, 0) "functools.partial", /* tp_name */ @@ -304,7 +351,7 @@ static PyTypeObject partial_type = { offsetof(partialobject, weakreflist), /* tp_weaklistoffset */ 0, /* tp_iter */ 0, /* tp_iternext */ - 0, /* tp_methods */ + partial_methods, /* tp_methods */ partial_memberlist, /* tp_members */ partial_getsetlist, /* tp_getset */ 0, /* tp_base */