bpo-31956: Add start and stop parameters to array.index() (GH-25059)

Co-Authored-By: Anders Lorentsen <Phaqui@gmail.com>

2 files changed, 61 insertions, 10 deletions

diff --git a/Modules/arraymodule.c b/Modules/arraymodule.c
index e7d5ab7..fb9ebbe 100644
--- a/Modules/arraymodule.c
+++ b/Modules/arraymodule.c
@@ -1136,18 +1136,32 @@ array_array_count(arrayobject *self, PyObject *v)
 array.array.index
 
     v: object
+    start: slice_index(accept={int}) = 0
+    stop: slice_index(accept={int}, c_default="PY_SSIZE_T_MAX") = sys.maxsize
     /
 
 Return index of first occurrence of v in the array.
+
+Raise ValueError if the value is not present.
 [clinic start generated code]*/
 
 static PyObject *
-array_array_index(arrayobject *self, PyObject *v)
-/*[clinic end generated code: output=d48498d325602167 input=cf619898c6649d08]*/
-{
-    Py_ssize_t i;
-
-    for (i = 0; i < Py_SIZE(self); i++) {
+array_array_index_impl(arrayobject *self, PyObject *v, Py_ssize_t start,
+                       Py_ssize_t stop)
+/*[clinic end generated code: output=c45e777880c99f52 input=089dff7baa7e5a7e]*/
+{
+    if (start < 0) {
+        start += Py_SIZE(self);
+        if (start < 0) {
+            start = 0;
+        }
+    }
+    if (stop < 0) {
+        stop += Py_SIZE(self);
+    }
+    // Use Py_SIZE() for every iteration in case the array is mutated
+    // during PyObject_RichCompareBool()
+    for (Py_ssize_t i = start; i < stop && i < Py_SIZE(self); i++) {
         PyObject *selfi;
         int cmp;
 
diff --git a/Modules/clinic/arraymodule.c.h b/Modules/clinic/arraymodule.c.h
index d0b70c4..d2513eb 100644
--- a/Modules/clinic/arraymodule.c.h
+++ b/Modules/clinic/arraymodule.c.h
@@ -39,13 +39,50 @@ PyDoc_STRVAR(array_array_count__doc__,
     {"count", (PyCFunction)array_array_count, METH_O, array_array_count__doc__},
 
 PyDoc_STRVAR(array_array_index__doc__,
-"index($self, v, /)\n"
+"index($self, v, start=0, stop=sys.maxsize, /)\n"
 "--\n"
 "\n"
-"Return index of first occurrence of v in the array.");
+"Return index of first occurrence of v in the array.\n"
+"\n"
+"Raise ValueError if the value is not present.");
 
 #define ARRAY_ARRAY_INDEX_METHODDEF    \
-    {"index", (PyCFunction)array_array_index, METH_O, array_array_index__doc__},
+    {"index", (PyCFunction)(void(*)(void))array_array_index, METH_FASTCALL, array_array_index__doc__},
+
+static PyObject *
+array_array_index_impl(arrayobject *self, PyObject *v, Py_ssize_t start,
+                       Py_ssize_t stop);
+
+static PyObject *
+array_array_index(arrayobject *self, PyObject *const *args, Py_ssize_t nargs)
+{
+    PyObject *return_value = NULL;
+    PyObject *v;
+    Py_ssize_t start = 0;
+    Py_ssize_t stop = PY_SSIZE_T_MAX;
+
+    if (!_PyArg_CheckPositional("index", nargs, 1, 3)) {
+        goto exit;
+    }
+    v = args[0];
+    if (nargs < 2) {
+        goto skip_optional;
+    }
+    if (!_PyEval_SliceIndexNotNone(args[1], &start)) {
+        goto exit;
+    }
+    if (nargs < 3) {
+        goto skip_optional;
+    }
+    if (!_PyEval_SliceIndexNotNone(args[2], &stop)) {
+        goto exit;
+    }
+skip_optional:
+    return_value = array_array_index_impl(self, v, start, stop);
+
+exit:
+    return return_value;
+}
 
 PyDoc_STRVAR(array_array_remove__doc__,
 "remove($self, v, /)\n"
@@ -535,4 +572,4 @@ PyDoc_STRVAR(array_arrayiterator___setstate____doc__,
 
 #define ARRAY_ARRAYITERATOR___SETSTATE___METHODDEF    \
     {"__setstate__", (PyCFunction)array_arrayiterator___setstate__, METH_O, array_arrayiterator___setstate____doc__},
-/*[clinic end generated code: output=a7f71a18b994c88f input=a9049054013a1b77]*/
+/*[clinic end generated code: output=376001addedc67ee input=a9049054013a1b77]*/