diff options
Diffstat (limited to 'Modules/_operator.c')
-rw-r--r-- | Modules/_operator.c | 42 |
1 files changed, 37 insertions, 5 deletions
diff --git a/Modules/_operator.c b/Modules/_operator.c index 3bf8c1276d..d6c6a18d81 100644 --- 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); |