#5228: add pickle support to functools.partial

This commit is contained in:
Jack Diederich 2009-03-31 23:46:48 +00:00
parent 840ac926c5
commit d60c29ed8b
4 changed files with 65 additions and 2 deletions

View File

@ -2,6 +2,7 @@ import functools
import unittest import unittest
from test import test_support from test import test_support
from weakref import proxy from weakref import proxy
import pickle
@staticmethod @staticmethod
def PythonPartial(func, *args, **keywords): def PythonPartial(func, *args, **keywords):
@ -19,6 +20,10 @@ def capture(*args, **kw):
"""capture all positional and keyword arguments""" """capture all positional and keyword arguments"""
return args, kw 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): class TestPartial(unittest.TestCase):
thetype = functools.partial thetype = functools.partial
@ -140,6 +145,12 @@ class TestPartial(unittest.TestCase):
join = self.thetype(''.join) join = self.thetype(''.join)
self.assertEqual(join(data), '0123456789') 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): class PartialSubclass(functools.partial):
pass pass
@ -147,11 +158,13 @@ class TestPartialSubclass(TestPartial):
thetype = PartialSubclass thetype = PartialSubclass
class TestPythonPartial(TestPartial): class TestPythonPartial(TestPartial):
thetype = PythonPartial thetype = PythonPartial
# the python version isn't picklable
def test_pickle(self): pass
class TestUpdateWrapper(unittest.TestCase): class TestUpdateWrapper(unittest.TestCase):
def check_wrapper(self, wrapper, wrapped, def check_wrapper(self, wrapper, wrapped,

View File

@ -168,6 +168,7 @@ Roger Dev
Raghuram Devarakonda Raghuram Devarakonda
Toby Dickenson Toby Dickenson
Mark Dickinson Mark Dickinson
Jack Diederich
Yves Dionne Yves Dionne
Daniel Dittmar Daniel Dittmar
Jaromir Dolecek Jaromir Dolecek

View File

@ -710,6 +710,8 @@ Extension Modules
- Issue #4396: The parser module now correctly validates the with statement. - Issue #4396: The parser module now correctly validates the with statement.
- Issue #5228: Make functools.partial objects can now be pickled.
Tests Tests
----- -----

View File

@ -274,6 +274,53 @@ static PyGetSetDef partial_getsetlist[] = {
{NULL} /* Sentinel */ {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 = { static PyTypeObject partial_type = {
PyVarObject_HEAD_INIT(NULL, 0) PyVarObject_HEAD_INIT(NULL, 0)
"functools.partial", /* tp_name */ "functools.partial", /* tp_name */
@ -304,7 +351,7 @@ static PyTypeObject partial_type = {
offsetof(partialobject, weakreflist), /* tp_weaklistoffset */ offsetof(partialobject, weakreflist), /* tp_weaklistoffset */
0, /* tp_iter */ 0, /* tp_iter */
0, /* tp_iternext */ 0, /* tp_iternext */
0, /* tp_methods */ partial_methods, /* tp_methods */
partial_memberlist, /* tp_members */ partial_memberlist, /* tp_members */
partial_getsetlist, /* tp_getset */ partial_getsetlist, /* tp_getset */
0, /* tp_base */ 0, /* tp_base */