operator.itemgetter() and operator.attrgetter() now support extraction
of multiple fields. This provides direct support for sorting by
multiple keys.
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 @@
\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 @@
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 @@
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 @@
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)
diff --git a/Misc/NEWS b/Misc/NEWS
index 36a21ce..3f4cbb1 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -47,6 +47,10 @@
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 @@
typedef struct {
PyObject_HEAD
+ int nitems;
PyObject *item;
} itemgetterobject;
@@ -266,9 +267,14 @@
{
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 @@
Py_INCREF(item);
ig->item = item;
+ ig->nitems = nitems;
PyObject_GC_Track(ig);
return (PyObject *)ig;
@@ -301,18 +308,40 @@
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 @@
typedef struct {
PyObject_HEAD
+ int nattrs;
PyObject *attr;
} attrgetterobject;
@@ -373,9 +403,14 @@
{
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 @@
Py_INCREF(attr);
ag->attr = attr;
+ ag->nattrs = nattrs;
PyObject_GC_Track(ag);
return (PyObject *)ag;
@@ -408,18 +444,40 @@
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)