bpo-35664: Optimize operator.itemgetter (pythonGH-11435)

rhettinger · web-flow · commit 2d53bed79c19 · 2019-01-07T09:38:41.000-07:00
diff --git a/Lib/test/test_operator.py b/Lib/test/test_operator.py
@@ -401,6 +401,19 @@ def __getitem__(self, name):
         self.assertEqual(operator.itemgetter(2,10,5)(data), ('2', '10', '5'))
         self.assertRaises(TypeError, operator.itemgetter(2, 'x', 5), data)
 
+        # interesting indices
+        t = tuple('abcde')
+        self.assertEqual(operator.itemgetter(-1)(t), 'e')
+        self.assertEqual(operator.itemgetter(slice(2, 4))(t), ('c', 'd'))
+
+        # interesting sequences
+        class T(tuple):
+            'Tuple subclass'
+            pass
+        self.assertEqual(operator.itemgetter(0)(T('abc')), 'a')
+        self.assertEqual(operator.itemgetter(0)(['a', 'b', 'c']), 'a')
+        self.assertEqual(operator.itemgetter(0)(range(100, 200)), 100)
+
     def test_methodcaller(self):
         operator = self.module
         self.assertRaises(TypeError, operator.methodcaller)
diff --git a/Misc/NEWS.d/next/Library/2019-01-04-22-18-25.bpo-35664.Z-Gyyj.rst b/Misc/NEWS.d/next/Library/2019-01-04-22-18-25.bpo-35664.Z-Gyyj.rst
@@ -0,0 +1,4 @@
+Improve operator.itemgetter() performance by 33% with optimized argument
+handling and with adding a fast path for the common case of a single
+non-negative integer index into a tuple (which is the typical use case in
+the standard library).
diff --git a/Modules/_operator.c b/Modules/_operator.c
@@ -937,6 +937,7 @@ typedef struct {
     PyObject_HEAD
     Py_ssize_t nitems;
     PyObject *item;
+    Py_ssize_t index; // -1 unless *item* is a single non-negative integer index
 } itemgetterobject;
 
 static PyTypeObject itemgetter_type;
@@ -948,6 +949,7 @@ itemgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
     itemgetterobject *ig;
     PyObject *item;
     Py_ssize_t nitems;
+    Py_ssize_t index;
 
     if (!_PyArg_NoKeywords("itemgetter", kwds))
         return NULL;
@@ -967,6 +969,21 @@ itemgetter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
     Py_INCREF(item);
     ig->item = item;
     ig->nitems = nitems;
+    ig->index = -1;
+    if (PyLong_CheckExact(item)) {
+        index = PyLong_AsSsize_t(item);
+        if (index < 0) {
+            /* If we get here, then either the index conversion failed
+             * due to being out of range, or the index was a negative
+             * integer.  Either way, we clear any possible exception
+             * and fall back to the slow path, where ig->index is -1.
+             */
+            PyErr_Clear();
+        }
+        else {
+            ig->index = index;
+        }
+    }
 
     PyObject_GC_Track(ig);
     return (PyObject *)ig;
@@ -993,12 +1010,27 @@ itemgetter_call(itemgetterobject *ig, PyObject *args, PyObject *kw)
     PyObject *obj, *result;
     Py_ssize_t i, nitems=ig->nitems;
 
-    if (!_PyArg_NoKeywords("itemgetter", kw))
-        return NULL;
-    if (!PyArg_UnpackTuple(args, "itemgetter", 1, 1, &obj))
-        return NULL;
-    if (nitems == 1)
+    assert(PyTuple_CheckExact(args));
+    if (kw == NULL && PyTuple_GET_SIZE(args) == 1) {
+        obj = PyTuple_GET_ITEM(args, 0);
+    }
+    else {
+        if (!_PyArg_NoKeywords("itemgetter", kw))
+            return NULL;
+        if (!PyArg_UnpackTuple(args, "itemgetter", 1, 1, &obj))
+            return NULL;
+    }
+    if (nitems == 1) {
+        if (ig->index >= 0
+            && PyTuple_CheckExact(obj)
+            && ig->index < PyTuple_GET_SIZE(obj))
+        {
+            result = PyTuple_GET_ITEM(obj, ig->index);
+            Py_INCREF(result);
+            return result;
+        }
         return PyObject_GetItem(obj, ig->item);
+    }
 
     assert(PyTuple_Check(ig->item));
     assert(PyTuple_GET_SIZE(ig->item) == nitems);
