diff -r 4990157343c6 Lib/operator.py --- a/Lib/operator.py Sat Nov 29 15:56:38 2014 +0100 +++ b/Lib/operator.py Sat Nov 29 16:04:45 2014 -0600 @@ -235,20 +235,25 @@ if not attrs: if not isinstance(attr, str): raise TypeError('attribute name must be a string') - names = attr.split('.') - def func(obj): - for name in names: - obj = getattr(obj, name) - return obj - self._call = func + self._names = attr.split('.') + self._getters = None else: - getters = tuple(map(attrgetter, (attr,) + attrs)) - def func(obj): - return tuple(getter(obj) for getter in getters) - self._call = func + self._getters = tuple(map(attrgetter, (attr,) + attrs)) def __call__(self, obj): - return self._call(obj) + if not self._getters: + for name in self._names: + obj = getattr(obj, name) + return obj + else: + return tuple(getter(obj) for getter in self._getters) + + def __repr__(self): + if not self._getters: + return '{}({})'.format(self.__class__.__name__, + '.'.join(self._names)) + attrs = ', '.join(('.'.join(getter._names) for getter in self._getters)) + return '{}({})'.format(self.__class__.__name__, attrs) class itemgetter: """ @@ -258,17 +263,23 @@ """ def __init__(self, item, *items): if not items: - def func(obj): - return obj[item] - self._call = func + self._item = item + self._nitems = 1 else: - items = (item,) + items - def func(obj): - return tuple(obj[i] for i in items) - self._call = func + self._items = (item,) + items + self._nitems = len(self._items) def __call__(self, obj): - return self._call(obj) + if self._nitems == 1: + return obj[self._item] + return tuple(obj[item] for item in self._items) + + def __repr__(self): + if self._nitems == 1: + return '{}({!r})'.format(self.__class__.__name__, self._item) + return '{}({})'.format(self.__class__.__name__, + ', '.join(map(repr, self._items))) + class methodcaller: """ @@ -290,6 +301,14 @@ def __call__(self, obj): return getattr(obj, self._name)(*self._args, **self._kwargs) + def __repr__(self): + posargs = ', '.join(map(repr, (self._name,) + self._args)) + kwdargs = ', '.join('{}={!r}'.format(k, v) + for k, v in self._kwargs.items()) + allargs = ', '.join(filter(None, (posargs, kwdargs))) + return '{}({})'.format(self.__class__.__name__, allargs) + + # In-place Operations *********************************************************# def iadd(a, b): diff -r 4990157343c6 Lib/test/test_operator.py --- a/Lib/test/test_operator.py Sat Nov 29 15:56:38 2014 +0100 +++ b/Lib/test/test_operator.py Sat Nov 29 16:04:45 2014 -0600 @@ -1,4 +1,6 @@ +import pickle import unittest +import sys from test import support @@ -35,6 +37,14 @@ class OperatorTestCase: + @classmethod + def setUpClass(cls): + sys.modules['operator'] = cls.module + + @classmethod + def tearDownClass(cls): + del sys.modules['operator'] + def test_lt(self): operator = self.module self.assertRaises(TypeError, operator.lt) @@ -358,6 +368,22 @@ f = operator.attrgetter('name', 'child.name', 'child.child.name') self.assertEqual(f(a), ('arthur', 'thomas', 'johnson')) + # Test pickle/unpickle + 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 +419,18 @@ self.assertEqual(operator.itemgetter(2,10,5)(data), ('2', '10', '5')) self.assertRaises(TypeError, operator.itemgetter(2, 'x', 5), data) + # Test pickle/unpickle + 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(f)) + 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 +454,27 @@ f = operator.methodcaller('baz', name='spam', self='eggs') self.assertEqual(f(a), ('spam', 'eggs')) + # Test pickle/unpickle + 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', self='eggs', name='spam') + pickled = pickle.loads(pickle.dumps(f)) + # Can't test repr consistently with keyword args + 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.assertRaises(KeyError, pickled, a) + def test_inplace(self): operator = self.module class C(object): diff -r 4990157343c6 Modules/_operator.c --- a/Modules/_operator.c Sat Nov 29 15:56:38 2014 +0100 +++ b/Modules/_operator.c Sat Nov 29 16:04:45 2014 -0600 @@ -485,6 +485,41 @@ 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) +{ + if (PyTuple_CheckExact(ig->item)) + return PyTuple_Pack(2, Py_TYPE(ig), ig->item); + return Py_BuildValue("O(O)", 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 +538,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 +556,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 +772,104 @@ 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); + return NULL; + } + } 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) +{ + Py_ssize_t i; + PyObject *attrsep = NULL; + PyObject *attrstrings = PyTuple_New(ag->nattrs); + if (attrstrings == NULL) + return NULL; + + for (i = 0; i < ag->nattrs; ++i) { + PyObject *attrstr; + PyObject *attr = PyTuple_GET_ITEM(ag->attr, i); + if (PyTuple_CheckExact(attr)) { + if (attrsep == NULL) { + attrsep = PyUnicode_FromString("."); + if (attrsep == NULL) { + Py_DECREF(attrstrings); + return NULL; + } + } + attrstr = PyUnicode_Join(attrsep, attr); + if (attrstr == NULL) { + Py_DECREF(attrstrings); + return NULL; + } + } else { + Py_INCREF(attr); + attrstr = attr; + } + PyTuple_SET_ITEM(attrstrings, i, attrstr); + } + return Py_BuildValue("ON", Py_TYPE(ag), attrstrings); +} + +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 +890,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 +908,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 +946,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 +965,8 @@ } mc->args = newargs; - name = PyTuple_GET_ITEM(args, 0); Py_INCREF(name); + PyUnicode_InternInPlace(&name); mc->name = name; Py_XINCREF(kwds); @@ -869,6 +1009,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 +1163,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 +1181,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 */