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

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

6 files changed, 87 insertions, 12 deletions

diff --git a/Doc/library/array.rst b/Doc/library/array.rst
index ff3ec6b..18a768d 100644
--- a/Doc/library/array.rst
+++ b/Doc/library/array.rst
@@ -178,11 +178,15 @@ The following data items and methods are also supported:
    array of some other type.
 
 
-.. method:: array.index(x)
+.. method:: array.index(x[, start[, stop]])
 
    Return the smallest *i* such that *i* is the index of the first occurrence of
-   *x* in the array.
+   *x* in the array.  The optional arguments *start* and *stop* can be
+   specified to search for *x* within a subsection of the array.  Raise
+   :exc:`ValueError` if *x* is not found.
 
+   .. versionchanged:: 3.10
+      Added optional *start* and *stop* parameters.
 
 .. method:: array.insert(i, x)
 
diff --git a/Doc/whatsnew/3.10.rst b/Doc/whatsnew/3.10.rst
index b84bcf9..eb42480 100644
--- a/Doc/whatsnew/3.10.rst
+++ b/Doc/whatsnew/3.10.rst
@@ -633,6 +633,13 @@ argparse
 Misleading phrase "optional arguments" was replaced with "options" in argparse help. Some tests might require adaptation if they rely on exact output match.
 (Contributed by Raymond Hettinger in :issue:`9694`.)
 
+array
+-----
+
+The :meth:`~array.array.index` method of :class:`array.array` now has
+optional *start* and *stop* parameters.
+(Contributed by Anders Lorentsen and Zackery Spytz in :issue:`31956`.)
+
 base64
 ------
 
diff --git a/Lib/test/test_array.py b/Lib/test/test_array.py
index 77a0c64..bdcd125 100644
--- a/Lib/test/test_array.py
+++ b/Lib/test/test_array.py
@@ -918,6 +918,17 @@ class BaseTest:
         self.assertRaises(ValueError, a.index, None)
         self.assertRaises(ValueError, a.index, self.outside)
 
+        a = array.array('i', [-2, -1, 0, 0, 1, 2])
+        self.assertEqual(a.index(0), 2)
+        self.assertEqual(a.index(0, 2), 2)
+        self.assertEqual(a.index(0, -4), 2)
+        self.assertEqual(a.index(-2, -10), 0)
+        self.assertEqual(a.index(0, 3), 3)
+        self.assertEqual(a.index(0, -3), 3)
+        self.assertEqual(a.index(0, 3, 4), 3)
+        self.assertEqual(a.index(0, -3, -2), 3)
+        self.assertRaises(ValueError, a.index, 2, 0, -10)
+
     def test_count(self):
         example = 2*self.example
         a = array.array(self.typecode, example)
diff --git a/Misc/NEWS.d/next/Library/2021-03-28-16-53-25.bpo-31956.Lt_67U.rst b/Misc/NEWS.d/next/Library/2021-03-28-16-53-25.bpo-31956.Lt_67U.rst
new file mode 100644
index 0000000..5285fd3
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2021-03-28-16-53-25.bpo-31956.Lt_67U.rst
@@ -0,0 +1,2 @@
+The :meth:`~array.array.index` method of :class:`array.array` now has
+optional *start* and *stop* parameters.
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]*/