diff -r 276515d2ceed Lib/test/test_operator.py --- a/Lib/test/test_operator.py Fri Nov 28 22:42:06 2014 +0100 +++ b/Lib/test/test_operator.py Fri Nov 28 21:48:13 2014 -0500 @@ -358,6 +358,23 @@ f = operator.attrgetter('name', 'child.name', 'child.child.name') self.assertEqual(f(a), ('arthur', 'thomas', 'johnson')) + # Test pickle/unpickle + import pickle + f = operator.attrgetter('name') + pickled = pickle.loads(pickle.dumps(f)) + self.assertEqual(repr(f), repr(pickled)) + self.assertEqual(f(a), pickled(a)) + + f = operator.attrgetter('child.child.name') + pickled = pickle.loads(pickle.dumps(f)) + self.assertEqual(repr(f), repr(pickled)) + self.assertEqual(f(a), pickled(a)) + + f = operator.attrgetter('name', 'child.name', 'child.child.name') + pickled = pickle.loads(pickle.dumps(f)) + self.assertEqual(repr(f), repr(pickled)) + self.assertEqual(f(a), pickled(a)) + def test_itemgetter(self): operator = self.module a = 'ABCDE' @@ -393,6 +410,19 @@ self.assertEqual(operator.itemgetter(2,10,5)(data), ('2', '10', '5')) self.assertRaises(TypeError, operator.itemgetter(2, 'x', 5), data) + # Test pickle/unpickle + import pickle + data = list(map(str, range(20))) + f = operator.itemgetter(2) + pickled = pickle.loads(pickle.dumps(f)) + self.assertEqual(repr(f), repr(pickled)) + self.assertEqual(f(data), pickled(data)) + + f = operator.itemgetter(5, 3, 7) + pickled = pickle.loads(pickle.dumps(multigetter)) + self.assertEqual(repr(f), repr(pickled)) + self.assertEqual(f(data), pickled(data)) + def test_methodcaller(self): operator = self.module self.assertRaises(TypeError, operator.methodcaller) @@ -416,6 +446,28 @@ f = operator.methodcaller('baz', name='spam', self='eggs') self.assertEqual(f(a), ('spam', 'eggs')) + # Test pickle/unpickle + import pickle + f = operator.methodcaller('bar') + pickled = pickle.loads(pickle.dumps(f)) + self.assertEqual(repr(f), repr(pickled)) + self.assertEqual(f(a), pickled(a)) + + f = operator.methodcaller('foo', 1, 2) + pickled = pickle.loads(pickle.dumps(f)) + self.assertEqual(repr(f), repr(pickled)) + self.assertEqual(f(a), pickled(a)) + + f = operator.methodcaller('baz', name='spam') + pickled = pickle.loads(pickle.dumps(f)) + self.assertEqual(repr(f), repr(pickled)) + self.assertEqual(f(a), pickled(a)) + + f = operator.methodcaller('baz', 1, 2, self='eggs') + pickled = pickle.loads(pickle.dumps(f)) + self.assertEqual(repr(f), repr(pickled)) + self.assertEqual(f(a), pickled(a)) + def test_inplace(self): operator = self.module class C(object): diff -r 276515d2ceed Modules/_operator.c --- a/Modules/_operator.c Fri Nov 28 22:42:06 2014 +0100 +++ b/Modules/_operator.c Fri Nov 28 21:48:13 2014 -0500 @@ -485,6 +485,39 @@ return result; } +static PyObject * +itemgetter_repr(itemgetterobject *ig) +{ + PyObject *repr; + const char *reprfmt; + + int status = Py_ReprEnter((PyObject *)ig); + if (status != 0) { + if (status < 0) + return NULL; + return PyUnicode_FromFormat("%s(...)", Py_TYPE(ig)->tp_name); + } + + reprfmt = ig->nitems == 1 ? "%s(%R)" : "%s%R"; + repr = PyUnicode_FromFormat(reprfmt, Py_TYPE(ig)->tp_name, ig->item); + Py_ReprLeave((PyObject *)ig); + return repr; +} + +static PyObject * +itemgetter_reduce(itemgetterobject *ig) +{ + return PyTuple_Pack(2, Py_TYPE(ig), ig->item); +} + +PyDoc_STRVAR(reduce_doc, "Return state information for pickling"); + +static PyMethodDef itemgetter_methods[] = { + {"__reduce__", (PyCFunction)itemgetter_reduce, METH_NOARGS, + reduce_doc}, + {NULL} +}; + PyDoc_STRVAR(itemgetter_doc, "itemgetter(item, ...) --> itemgetter object\n\ \n\ @@ -503,7 +536,7 @@ 0, /* tp_getattr */ 0, /* tp_setattr */ 0, /* tp_reserved */ - 0, /* tp_repr */ + (reprfunc)itemgetter_repr, /* tp_repr */ 0, /* tp_as_number */ 0, /* tp_as_sequence */ 0, /* tp_as_mapping */ @@ -521,7 +554,7 @@ 0, /* tp_weaklistoffset */ 0, /* tp_iter */ 0, /* tp_iternext */ - 0, /* tp_methods */ + itemgetter_methods, /* tp_methods */ 0, /* tp_members */ 0, /* tp_getset */ 0, /* tp_base */ @@ -737,6 +770,75 @@ return result; } +static PyObject * +attrgetter_repr(attrgetterobject *ag) +{ + PyObject *attrstrings, *repr, *attrsep; + Py_ssize_t i; + + int status = Py_ReprEnter((PyObject *)ag); + if (status != 0) { + if (status < 0) + return NULL; + return PyUnicode_FromFormat("%s(...)", Py_TYPE(ag)->tp_name); + } + + if (ag->nattrs == 1) { + repr = PyUnicode_FromFormat("%s(%R)", Py_TYPE(ag)->tp_name, PyTuple_GET_ITEM(ag->attr, 0)); + Py_ReprLeave((PyObject *)ag); + return repr; + } + + attrstrings = PyTuple_New(ag->nattrs); + if (attrstrings == NULL) { + Py_ReprLeave((PyObject *)ag); + return NULL; + } + + attrsep = PyUnicode_FromString("."); + if (attrsep == NULL) { + Py_DECREF(attrstrings); + Py_ReprLeave((PyObject *)ag); + return NULL; + } + + for (i = 0; i < ag->nattrs; ++i) { + PyObject *attrstr; + PyObject *attr = PyTuple_GET_ITEM(ag->attr, i); + + if (PyTuple_CheckExact(attr)) { + attrstr = PyUnicode_Join(attrsep, attr); + if (attrstr == NULL) { + Py_DECREF(attrstrings); + Py_DECREF(attrsep); + Py_ReprLeave((PyObject *)ag); + } + } else { + Py_INCREF(attr); + attrstr = attr; + } + PyTuple_SET_ITEM(attrstrings, i, attrstr); + } + Py_DECREF(attrsep); + + repr = PyUnicode_FromFormat("%s%R", Py_TYPE(ag)->tp_name, attrstrings); + Py_DECREF(attrstrings); + Py_ReprLeave((PyObject *)ag); + return repr; +} + +static PyObject * +attrgetter_reduce(attrgetterobject *ag) +{ + return PyTuple_Pack(2, Py_TYPE(ag), ag->attr); +} + +static PyMethodDef attrgetter_methods[] = { + {"__reduce__", (PyCFunction)attrgetter_reduce, METH_NOARGS, + reduce_doc}, + {NULL} +}; + PyDoc_STRVAR(attrgetter_doc, "attrgetter(attr, ...) --> attrgetter object\n\ \n\ @@ -757,7 +859,7 @@ 0, /* tp_getattr */ 0, /* tp_setattr */ 0, /* tp_reserved */ - 0, /* tp_repr */ + (reprfunc)attrgetter_repr, /* tp_repr */ 0, /* tp_as_number */ 0, /* tp_as_sequence */ 0, /* tp_as_mapping */ @@ -775,7 +877,7 @@ 0, /* tp_weaklistoffset */ 0, /* tp_iter */ 0, /* tp_iternext */ - 0, /* tp_methods */ + attrgetter_methods, /* tp_methods */ 0, /* tp_members */ 0, /* tp_getset */ 0, /* tp_base */ @@ -813,6 +915,13 @@ return NULL; } + name = PyTuple_GET_ITEM(args, 0); + if (!PyUnicode_Check(name)) { + PyErr_SetString(PyExc_TypeError, + "method name must be a string"); + return NULL; + } + /* create methodcallerobject structure */ mc = PyObject_GC_New(methodcallerobject, &methodcaller_type); if (mc == NULL) @@ -825,8 +934,8 @@ } mc->args = newargs; - name = PyTuple_GET_ITEM(args, 0); Py_INCREF(name); + PyUnicode_InternInPlace(&name); mc->name = name; Py_XINCREF(kwds); @@ -869,6 +978,141 @@ return result; } +static PyObject * +methodcaller_repr(methodcallerobject *mc) +{ + PyObject *argreprs, *repr, *sep, *joinedargreprs; + Py_ssize_t numtotalargs, numposargs, numkwdargs, i; + int status = Py_ReprEnter((PyObject *)mc); + if (status != 0) { + if (status < 0) + return NULL; + return PyUnicode_FromFormat("%s(...)", Py_TYPE(mc)->tp_name); + } + + if (mc->kwds != NULL) { + numkwdargs = PyDict_Size(mc->kwds); + if (numkwdargs < 0) { + Py_ReprLeave((PyObject *)mc); + return NULL; + } + } else { + numkwdargs = 0; + } + + numposargs = PyTuple_GET_SIZE(mc->args); + numtotalargs = numposargs + numkwdargs; + + if (numtotalargs == 0) { + repr = PyUnicode_FromFormat("%s(%R)", Py_TYPE(mc)->tp_name, mc->name); + Py_ReprLeave((PyObject *)mc); + return repr; + } + + argreprs = PyTuple_New(numtotalargs); + if (argreprs == NULL) + return NULL; + + for (i = 0; i < numposargs; ++i) { + PyObject *onerepr = PyObject_Repr(PyTuple_GET_ITEM(mc->args, i)); + if (onerepr == NULL) { + Py_DECREF(argreprs); + Py_ReprLeave((PyObject *)mc); + return NULL; + } + PyTuple_SET_ITEM(argreprs, i, onerepr); + } + + if (mc->kwds != NULL) { + PyObject *key, *value; + Py_ssize_t pos = 0; + while (PyDict_Next(mc->kwds, &pos, &key, &value)) { + PyObject *onerepr = PyUnicode_FromFormat("%U=%R", key, value); + if (onerepr == NULL) { + Py_DECREF(argreprs); + Py_ReprLeave((PyObject *)mc); + return NULL; + } + PyTuple_SET_ITEM(argreprs, i, onerepr); + ++i; + } + } + assert(i == numtotalargs); + + sep = PyUnicode_FromString(", "); + if (sep == NULL) { + Py_DECREF(argreprs); + Py_ReprLeave((PyObject *)mc); + return NULL; + } + + joinedargreprs = PyUnicode_Join(sep, argreprs); + Py_DECREF(sep); + Py_DECREF(argreprs); + if (joinedargreprs == NULL) { + Py_ReprLeave((PyObject *)mc); + return NULL; + } + + repr = PyUnicode_FromFormat("%s(%R, %U)", Py_TYPE(mc)->tp_name, mc->name, joinedargreprs); + Py_DECREF(joinedargreprs); + Py_ReprLeave((PyObject *)mc); + return repr; +} + +static PyObject * +methodcaller_reduce(methodcallerobject *mc) +{ + PyObject *newargs; + Py_ssize_t i; + Py_ssize_t callargcount = PyTuple_GET_SIZE(mc->args); + + newargs = PyTuple_New(1 + callargcount); + if (newargs == NULL) + return NULL; + + Py_INCREF(mc->name); + PyTuple_SET_ITEM(newargs, 0, mc->name); + for (i = 0; i < callargcount; ++i) { + PyObject *arg = PyTuple_GET_ITEM(mc->args, i); + Py_INCREF(arg); + PyTuple_SET_ITEM(newargs, i + 1, arg); + } + + if (mc->kwds && PyDict_Size(mc->kwds) > 0) { + return Py_BuildValue("ONO", Py_TYPE(mc), newargs, mc->kwds); + } + return Py_BuildValue("ON", Py_TYPE(mc), newargs); +} + +static PyObject * +methodcaller_setstate(methodcallerobject *mc, PyObject *state) +{ + if (state != NULL && !PyDict_Check(state)) { + PyErr_SetString(PyExc_TypeError, + "__setstate__ requires a dict"); + return NULL; + } + + Py_CLEAR(mc->kwds); + + if (state != NULL) { + Py_INCREF(state); + mc->kwds = state; + } + + Py_RETURN_NONE; +} + +PyDoc_STRVAR(setstate_doc, "Set state information for unpickling."); + +static PyMethodDef methodcaller_methods[] = { + {"__reduce__", (PyCFunction)methodcaller_reduce, METH_NOARGS, + reduce_doc}, + {"__setstate__", (PyCFunction)methodcaller_setstate, METH_O, + setstate_doc}, + {NULL} +}; PyDoc_STRVAR(methodcaller_doc, "methodcaller(name, ...) --> methodcaller object\n\ \n\ @@ -888,7 +1132,7 @@ 0, /* tp_getattr */ 0, /* tp_setattr */ 0, /* tp_reserved */ - 0, /* tp_repr */ + (reprfunc)methodcaller_repr, /* tp_repr */ 0, /* tp_as_number */ 0, /* tp_as_sequence */ 0, /* tp_as_mapping */ @@ -906,7 +1150,7 @@ 0, /* tp_weaklistoffset */ 0, /* tp_iter */ 0, /* tp_iternext */ - 0, /* tp_methods */ + methodcaller_methods, /* tp_methods */ 0, /* tp_members */ 0, /* tp_getset */ 0, /* tp_base */