]> granicus.if.org Git - python/commitdiff
operator.itemgetter() and operator.attrgetter() now support extraction
authorRaymond Hettinger <python@rcn.com>
Wed, 9 Mar 2005 16:38:48 +0000 (16:38 +0000)
committerRaymond Hettinger <python@rcn.com>
Wed, 9 Mar 2005 16:38:48 +0000 (16:38 +0000)
of multiple fields.  This provides direct support for sorting by
multiple keys.

Doc/lib/liboperator.tex
Lib/test/test_operator.py
Misc/NEWS
Modules/operator.c

index 52c0e932a162ba6c9b9f94a8765930c6a8476f93..d3c974a0faf604da7027e6033b9faa9c3bc6c6c1 100644 (file)
@@ -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)
index f699cd79911bea72d1db88e05c61888b36ff9a58..725b2d947750f6f979a2dc82974a0d68b1a80e83 100644 (file)
@@ -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)
index 36a21ce3dc496bbcec00c83e5629c2db1b2b0e97..3f4cbb1250650cee985f54a1cdc182caa9f50a9e 100644 (file)
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -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
index 468440e7aae43bf2da63acba8c517157acccb2f0..938c5e53df2704b907749db8e8f1562d1a6584bf 100644 (file)
@@ -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)