As discussed on python-dev, added two extractor functions to the

operator module.

4 files changed, 302 insertions, 2 deletions

diff --git a/Doc/lib/liboperator.tex b/Doc/lib/liboperator.tex
index e3aae60706..ab68d0e01f 100644
--- a/Doc/lib/liboperator.tex
+++ b/Doc/lib/liboperator.tex
@@ -300,6 +300,39 @@ Example: Build a dictionary that maps the ordinals from \code{0} to
 \end{verbatim}
 
 
+The \module{operator} module also defines tools for generalized attribute
+and item lookups.  These are useful for making fast field extractors
+as arguments for \function{map()}, \method{list.sort()},
+\method{itertools.groupby()}, or other functions that expect a
+function argument.
+
+\begin{funcdesc}{attrgetter}{attr}
+Return a callable object that fetches \var{attr} from its operand.
+After, \samp{f=attrgetter('name')}, the call \samp{f(b)} returns
+\samp{b.name}.
+\versionadded{2.4}
+\end{funcdesc}
+    
+\begin{funcdesc}{itemgetter}{item}
+Return a callable object that fetches \var{item} from its operand.
+After, \samp{f=itemgetter(2)}, the call \samp{f(b)} returns
+\samp{b[2]}.
+\versionadded{2.4}
+\end{funcdesc}
+
+Examples:
+                
+\begin{verbatim}
+>>> from operator import *
+>>> inventory = [('apple', 3), ('banana', 2), ('pear', 5), ('orange', 1)]
+>>> getcount = itemgetter(1)
+>>> map(getcount, inventory)
+[3, 2, 5, 1]
+>>> list.sorted(inventory, key=getcount)
+[('orange', 1), ('banana', 2), ('apple', 3), ('pear', 5)]
+\end{verbatim}
+                
+
 \subsection{Mapping Operators to Functions \label{operator-map}}
 
 This table shows how abstract operations correspond to operator
diff --git a/Lib/test/test_operator.py b/Lib/test/test_operator.py
index 422a3cb74a..e3a67f0d6d 100644
--- a/Lib/test/test_operator.py
+++ b/Lib/test/test_operator.py
@@ -227,6 +227,45 @@ class OperatorTestCase(unittest.TestCase):
         self.failIf(operator.is_not(a, b))
         self.failUnless(operator.is_not(a,c))
 
+    def test_attrgetter(self):
+        class A:
+            pass
+        a = A()
+        a.name = 'arthur'
+        f = operator.attrgetter('name')
+        self.assertEqual(f(a), 'arthur')
+        f = operator.attrgetter('rank')
+        self.assertRaises(AttributeError, f, a)
+        f = operator.attrgetter(2)
+        self.assertRaises(TypeError, f, a)
+        self.assertRaises(TypeError, operator.attrgetter)
+        self.assertRaises(TypeError, operator.attrgetter, 1, 2)
+
+    def test_itemgetter(self):
+        a = 'ABCDE'
+        f = operator.itemgetter(2)
+        self.assertEqual(f(a), 'C')
+        f = operator.itemgetter(10)
+        self.assertRaises(IndexError, f, a)
+
+        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')
+        self.assertEqual(f(d), 'val')
+        f = operator.itemgetter('nonkey')
+        self.assertRaises(KeyError, f, d)
+
+        # example used in the docs
+        inventory = [('apple', 3), ('banana', 2), ('pear', 5), ('orange', 1)]
+        getcount = operator.itemgetter(1)
+        self.assertEqual(map(getcount, inventory), [3, 2, 5, 1])
+        self.assertEqual(list.sorted(inventory, key=getcount),
+            [('orange', 1), ('banana', 2), ('apple', 3), ('pear', 5)])
+
 def test_main():
     test_support.run_unittest(OperatorTestCase)
 
diff --git a/Misc/NEWS b/Misc/NEWS
index ce9d779e2d..8f1c1a1a89 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -104,6 +104,11 @@ Core and builtins
 Extension modules
 -----------------
 
+- The operator module has two new functions, attrgetter() and
+  itemgetter() which are useful for creating fast data extractor
+  functions for map(), list.sort(), itertools.groupby(), and
+  other functions that expect a function argument.
+
 - socket.SHUT_{RD,WR,RDWR} was added.
 
 - os.getsid was added.
diff --git a/Modules/operator.c b/Modules/operator.c
index 7638fb80a1..d8e2a54602 100644
--- a/Modules/operator.c
+++ b/Modules/operator.c
@@ -252,13 +252,236 @@ spam2(ge,__ge__, "ge(a, b) -- Same as a>=b.")
 
 };
 
+/* itemgetter object **********************************************************/
 
+typedef struct {
+	PyObject_HEAD
+	PyObject *item;
+} itemgetterobject;
+
+static PyTypeObject itemgetter_type;
+
+static PyObject *
+itemgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
+{
+	itemgetterobject *ig;
+	PyObject *item;
+
+	if (!PyArg_UnpackTuple(args, "itemgetter", 1, 1, &item))
+		return NULL;
+
+	/* create itemgetterobject structure */
+	ig = PyObject_GC_New(itemgetterobject, &itemgetter_type);
+	if (ig == NULL) 
+		return NULL;	
+	
+	Py_INCREF(item);
+	ig->item = item;
+
+	PyObject_GC_Track(ig);
+	return (PyObject *)ig;
+}
+
+static void
+itemgetter_dealloc(itemgetterobject *ig)
+{
+	PyObject_GC_UnTrack(ig);
+	Py_XDECREF(ig->item);
+	PyObject_GC_Del(ig);
+}
+
+static int
+itemgetter_traverse(itemgetterobject *ig, visitproc visit, void *arg)
+{
+	if (ig->item)
+		return visit(ig->item, arg);
+	return 0;
+}
+
+static PyObject *
+itemgetter_call(itemgetterobject *ig, PyObject *args, PyObject *kw)
+{
+	PyObject * obj;
+
+	if (!PyArg_UnpackTuple(args, "itemgetter", 1, 1, &obj))
+		return NULL;
+	return PyObject_GetItem(obj, ig->item);
+}
+
+PyDoc_STRVAR(itemgetter_doc,
+"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].");
+
+static PyTypeObject itemgetter_type = {
+	PyObject_HEAD_INIT(NULL)
+	0,				/* ob_size */
+	"itertools.itemgetter",		/* tp_name */
+	sizeof(itemgetterobject),	/* tp_basicsize */
+	0,				/* tp_itemsize */
+	/* methods */
+	(destructor)itemgetter_dealloc,	/* tp_dealloc */
+	0,				/* tp_print */
+	0,				/* tp_getattr */
+	0,				/* tp_setattr */
+	0,				/* tp_compare */
+	0,				/* tp_repr */
+	0,				/* tp_as_number */
+	0,				/* tp_as_sequence */
+	0,				/* tp_as_mapping */
+	0,				/* tp_hash */
+	(ternaryfunc)itemgetter_call,	/* tp_call */
+	0,				/* tp_str */
+	PyObject_GenericGetAttr,	/* tp_getattro */
+	0,				/* tp_setattro */
+	0,				/* tp_as_buffer */
+	Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,	/* tp_flags */
+	itemgetter_doc,			/* tp_doc */
+	(traverseproc)itemgetter_traverse,	/* tp_traverse */
+	0,				/* tp_clear */
+	0,				/* tp_richcompare */
+	0,				/* tp_weaklistoffset */
+	0,				/* tp_iter */
+	0,				/* tp_iternext */
+	0,				/* tp_methods */
+	0,				/* tp_members */
+	0,				/* tp_getset */
+	0,				/* tp_base */
+	0,				/* tp_dict */
+	0,				/* tp_descr_get */
+	0,				/* tp_descr_set */
+	0,				/* tp_dictoffset */
+	0,				/* tp_init */
+	0,				/* tp_alloc */
+	itemgetter_new,			/* tp_new */
+	0,				/* tp_free */
+};
+
+
+/* attrgetter object **********************************************************/
+
+typedef struct {
+	PyObject_HEAD
+	PyObject *attr;
+} attrgetterobject;
+
+static PyTypeObject attrgetter_type;
+
+static PyObject *
+attrgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
+{
+	attrgetterobject *ag;
+	PyObject *attr;
+
+	if (!PyArg_UnpackTuple(args, "attrgetter", 1, 1, &attr))
+		return NULL;
+
+	/* create attrgetterobject structure */
+	ag = PyObject_GC_New(attrgetterobject, &attrgetter_type);
+	if (ag == NULL) 
+		return NULL;	
+	
+	Py_INCREF(attr);
+	ag->attr = attr;
+
+	PyObject_GC_Track(ag);
+	return (PyObject *)ag;
+}
+
+static void
+attrgetter_dealloc(attrgetterobject *ag)
+{
+	PyObject_GC_UnTrack(ag);
+	Py_XDECREF(ag->attr);
+	PyObject_GC_Del(ag);
+}
+
+static int
+attrgetter_traverse(attrgetterobject *ag, visitproc visit, void *arg)
+{
+	if (ag->attr)
+		return visit(ag->attr, arg);
+	return 0;
+}
+
+static PyObject *
+attrgetter_call(attrgetterobject *ag, PyObject *args, PyObject *kw)
+{
+	PyObject * obj;
+
+	if (!PyArg_UnpackTuple(args, "attrgetter", 1, 1, &obj))
+		return NULL;
+	return PyObject_GetAttr(obj, ag->attr);
+}
+
+PyDoc_STRVAR(attrgetter_doc,
+"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.");
+
+static PyTypeObject attrgetter_type = {
+	PyObject_HEAD_INIT(NULL)
+	0,				/* ob_size */
+	"itertools.attrgetter",		/* tp_name */
+	sizeof(attrgetterobject),	/* tp_basicsize */
+	0,				/* tp_itemsize */
+	/* methods */
+	(destructor)attrgetter_dealloc,	/* tp_dealloc */
+	0,				/* tp_print */
+	0,				/* tp_getattr */
+	0,				/* tp_setattr */
+	0,				/* tp_compare */
+	0,				/* tp_repr */
+	0,				/* tp_as_number */
+	0,				/* tp_as_sequence */
+	0,				/* tp_as_mapping */
+	0,				/* tp_hash */
+	(ternaryfunc)attrgetter_call,	/* tp_call */
+	0,				/* tp_str */
+	PyObject_GenericGetAttr,	/* tp_getattro */
+	0,				/* tp_setattro */
+	0,				/* tp_as_buffer */
+	Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,	/* tp_flags */
+	attrgetter_doc,			/* tp_doc */
+	(traverseproc)attrgetter_traverse,	/* tp_traverse */
+	0,				/* tp_clear */
+	0,				/* tp_richcompare */
+	0,				/* tp_weaklistoffset */
+	0,				/* tp_iter */
+	0,				/* tp_iternext */
+	0,				/* tp_methods */
+	0,				/* tp_members */
+	0,				/* tp_getset */
+	0,				/* tp_base */
+	0,				/* tp_dict */
+	0,				/* tp_descr_get */
+	0,				/* tp_descr_set */
+	0,				/* tp_dictoffset */
+	0,				/* tp_init */
+	0,				/* tp_alloc */
+	attrgetter_new,			/* tp_new */
+	0,				/* tp_free */
+};
 /* Initialization function for the module (*must* be called initoperator) */
 
 PyMODINIT_FUNC
 initoperator(void)
 {
-        /* Create the module and add the functions */
-        Py_InitModule4("operator", operator_methods, operator_doc,
+	PyObject *m;
+        
+	/* Create the module and add the functions */
+        m = Py_InitModule4("operator", operator_methods, operator_doc,
 		       (PyObject*)NULL, PYTHON_API_VERSION);
+
+	if (PyType_Ready(&itemgetter_type) < 0)
+		return;
+	Py_INCREF(&itemgetter_type);
+	PyModule_AddObject(m, "itemgetter", (PyObject *)&itemgetter_type);
+
+	if (PyType_Ready(&attrgetter_type) < 0)
+		return;
+	Py_INCREF(&attrgetter_type);
+	PyModule_AddObject(m, "attrgetter", (PyObject *)&attrgetter_type);
 }