diff options
| -rw-r--r-- | Doc/lib/liboperator.tex | 15 | ||||
| -rw-r--r-- | Lib/test/test_operator.py | 36 | ||||
| -rw-r--r-- | Misc/NEWS | 4 | ||||
| -rw-r--r-- | Modules/operator.c | 86 |
4 files changed, 118 insertions, 23 deletions
diff --git a/Doc/lib/liboperator.tex b/Doc/lib/liboperator.tex index 52c0e93..d3c974a 100644 --- a/Doc/lib/liboperator.tex +++ b/Doc/lib/liboperator.tex @@ -306,24 +306,31 @@ as arguments for \function{map()}, \function{sorted()}, \method{itertools.groupby()}, or other functions that expect a function argument. -\begin{funcdesc}{attrgetter}{attr} +\begin{funcdesc}{attrgetter}{attr\optional{, args...}} Return a callable object that fetches \var{attr} from its operand. +If more than one attribute is requested, returns a tuple of attributes. After, \samp{f=attrgetter('name')}, the call \samp{f(b)} returns -\samp{b.name}. +\samp{b.name}. After, \samp{f=attrgetter('name', 'date')}, the call +\samp{f(b)} returns \samp{(b.name, b.date)}. \versionadded{2.4} +\versionchanged[Added support for multiple attributes]{2.5} \end{funcdesc} -\begin{funcdesc}{itemgetter}{item} +\begin{funcdesc}{itemgetter}{item\optional{, args...}} Return a callable object that fetches \var{item} from its operand. +If more than one item is requested, returns a tuple of items. After, \samp{f=itemgetter(2)}, the call \samp{f(b)} returns \samp{b[2]}. +After, \samp{f=itemgetter(2,5,3)}, the call \samp{f(b)} returns +\samp{(b[2], b[5], b[3])}. \versionadded{2.4} +\versionchanged[Added support for multiple item extraction]{2.5} \end{funcdesc} Examples: \begin{verbatim} ->>> from operator import * +>>> from operator import itemgetter >>> inventory = [('apple', 3), ('banana', 2), ('pear', 5), ('orange', 1)] >>> getcount = itemgetter(1) >>> map(getcount, inventory) diff --git a/Lib/test/test_operator.py b/Lib/test/test_operator.py index f699cd7..725b2d9 100644 --- a/Lib/test/test_operator.py +++ b/Lib/test/test_operator.py @@ -324,7 +324,14 @@ class OperatorTestCase(unittest.TestCase): f = operator.attrgetter(2) self.assertRaises(TypeError, f, a) self.assertRaises(TypeError, operator.attrgetter) - self.assertRaises(TypeError, operator.attrgetter, 1, 2) + + # multiple gets + record = A() + record.x = 'X' + record.y = 'Y' + record.z = 'Z' + self.assertEqual(operator.attrgetter('x','z','y')(record), ('X', 'Z', 'Y')) + self.assertRaises(TypeError, operator.attrgetter('x', (), 'y'), record) class C(object): def __getattr(self, name): @@ -346,7 +353,6 @@ class OperatorTestCase(unittest.TestCase): f = operator.itemgetter('name') self.assertRaises(TypeError, f, a) self.assertRaises(TypeError, operator.itemgetter) - self.assertRaises(TypeError, operator.itemgetter, 1, 2) d = dict(key='val') f = operator.itemgetter('key') @@ -361,9 +367,29 @@ class OperatorTestCase(unittest.TestCase): self.assertEqual(sorted(inventory, key=getcount), [('orange', 1), ('banana', 2), ('apple', 3), ('pear', 5)]) -def test_main(): - test_support.run_unittest(OperatorTestCase) + # multiple gets + data = map(str, range(20)) + self.assertEqual(operator.itemgetter(2,10,5)(data), ('2', '10', '5')) + self.assertRaises(TypeError, operator.itemgetter(2, 'x', 5), data) + + +def test_main(verbose=None): + import sys + test_classes = ( + OperatorTestCase, + ) + + test_support.run_unittest(*test_classes) + # verify reference counting + if verbose and hasattr(sys, "gettotalrefcount"): + import gc + counts = [None] * 5 + for i in xrange(len(counts)): + test_support.run_unittest(*test_classes) + gc.collect() + counts[i] = sys.gettotalrefcount() + print counts if __name__ == "__main__": - test_main() + test_main(verbose=True) @@ -47,6 +47,10 @@ Core and builtins Extension Modules ----------------- +- operator.itemgetter() and operator.attrgetter() now support retrieving + multiple fields. This provides direct support for sorting on multiple + keys (primary, secondary, etc). + - os.access now supports Unicode path names on non-Win32 systems. - Patches #925152, #1118602: Avoid reading after the end of the buffer diff --git a/Modules/operator.c b/Modules/operator.c index 468440e..938c5e5 100644 --- a/Modules/operator.c +++ b/Modules/operator.c @@ -256,6 +256,7 @@ spam2(ge,__ge__, "ge(a, b) -- Same as a>=b.") typedef struct { PyObject_HEAD + int nitems; PyObject *item; } itemgetterobject; @@ -266,9 +267,14 @@ itemgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds) { itemgetterobject *ig; PyObject *item; + int nitems; - if (!PyArg_UnpackTuple(args, "itemgetter", 1, 1, &item)) - return NULL; + nitems = PyTuple_GET_SIZE(args); + if (nitems <= 1) { + if (!PyArg_UnpackTuple(args, "itemgetter", 1, 1, &item)) + return NULL; + } else + item = args; /* create itemgetterobject structure */ ig = PyObject_GC_New(itemgetterobject, &itemgetter_type); @@ -277,6 +283,7 @@ itemgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds) Py_INCREF(item); ig->item = item; + ig->nitems = nitems; PyObject_GC_Track(ig); return (PyObject *)ig; @@ -301,18 +308,40 @@ itemgetter_traverse(itemgetterobject *ig, visitproc visit, void *arg) static PyObject * itemgetter_call(itemgetterobject *ig, PyObject *args, PyObject *kw) { - PyObject * obj; + PyObject *obj, *result; + int i, nitems=ig->nitems; if (!PyArg_UnpackTuple(args, "itemgetter", 1, 1, &obj)) return NULL; - return PyObject_GetItem(obj, ig->item); + if (nitems == 1) + return PyObject_GetItem(obj, ig->item); + + assert(PyTuple_Check(ig->item)); + assert(PyTuple_GET_SIZE(ig->item) == nitems); + + result = PyTuple_New(nitems); + if (result == NULL) + return NULL; + + for (i=0 ; i < nitems ; i++) { + PyObject *item, *val; + item = PyTuple_GET_ITEM(ig->item, i); + val = PyObject_GetItem(obj, item); + if (val == NULL) { + Py_DECREF(result); + return NULL; + } + PyTuple_SET_ITEM(result, i, val); + } + return result; } PyDoc_STRVAR(itemgetter_doc, -"itemgetter(item) --> itemgetter object\n\ +"itemgetter(item, ...) --> itemgetter object\n\ \n\ -Return a callable object that fetches the given item from its operand.\n\ -After, f=itemgetter(2), the call f(b) returns b[2]."); +Return a callable object that fetches the given item(s) from its operand.\n\ +After, f=itemgetter(2), the call f(r) returns r[2].\n\ +After, g=itemgetter(2,5,3), the call g(r) returns (r[2], r[5], r[3])"); static PyTypeObject itemgetter_type = { PyObject_HEAD_INIT(NULL) @@ -363,6 +392,7 @@ static PyTypeObject itemgetter_type = { typedef struct { PyObject_HEAD + int nattrs; PyObject *attr; } attrgetterobject; @@ -373,9 +403,14 @@ attrgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds) { attrgetterobject *ag; PyObject *attr; + int nattrs; - if (!PyArg_UnpackTuple(args, "attrgetter", 1, 1, &attr)) - return NULL; + nattrs = PyTuple_GET_SIZE(args); + if (nattrs <= 1) { + if (!PyArg_UnpackTuple(args, "attrgetter", 1, 1, &attr)) + return NULL; + } else + attr = args; /* create attrgetterobject structure */ ag = PyObject_GC_New(attrgetterobject, &attrgetter_type); @@ -384,6 +419,7 @@ attrgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds) Py_INCREF(attr); ag->attr = attr; + ag->nattrs = nattrs; PyObject_GC_Track(ag); return (PyObject *)ag; @@ -408,18 +444,40 @@ attrgetter_traverse(attrgetterobject *ag, visitproc visit, void *arg) static PyObject * attrgetter_call(attrgetterobject *ag, PyObject *args, PyObject *kw) { - PyObject * obj; + PyObject *obj, *result; + int i, nattrs=ag->nattrs; if (!PyArg_UnpackTuple(args, "attrgetter", 1, 1, &obj)) return NULL; - return PyObject_GetAttr(obj, ag->attr); + if (ag->nattrs == 1) + return PyObject_GetAttr(obj, ag->attr); + + assert(PyTuple_Check(ag->attr)); + assert(PyTuple_GET_SIZE(ag->attr) == nattrs); + + result = PyTuple_New(nattrs); + if (result == NULL) + return NULL; + + for (i=0 ; i < nattrs ; i++) { + PyObject *attr, *val; + attr = PyTuple_GET_ITEM(ag->attr, i); + val = PyObject_GetAttr(obj, attr); + if (val == NULL) { + Py_DECREF(result); + return NULL; + } + PyTuple_SET_ITEM(result, i, val); + } + return result; } PyDoc_STRVAR(attrgetter_doc, -"attrgetter(attr) --> attrgetter object\n\ +"attrgetter(attr, ...) --> attrgetter object\n\ \n\ -Return a callable object that fetches the given attribute from its operand.\n\ -After, f=attrgetter('name'), the call f(b) returns b.name."); +Return a callable object that fetches the given attribute(s) from its operand.\n\ +After, f=attrgetter('name'), the call f(r) returns r.name.\n\ +After, g=attrgetter('name', 'date'), the call g(r) returns (r.name, r.date)."); static PyTypeObject attrgetter_type = { PyObject_HEAD_INIT(NULL) |