diff --git a/Lib/test/test_array.py b/Lib/test/test_array.py index e9218f3dd6..b77c368047 100644 --- a/Lib/test/test_array.py +++ b/Lib/test/test_array.py @@ -871,6 +871,25 @@ class BaseTest: self.assertRaises(ValueError, a.index, None) self.assertRaises(ValueError, a.index, self.outside) + for i, x in enumerate(example): + idx = example.index(x) + self.assertEqual(a.index(x, idx), example.index(x)) + + 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) + # can't run this tests because of OverflowError: They don't fit in + # Py_ssize_t - is this ok ? + #self.assertEqual(a.index(0, -4*sys.maxsize, 4*sys.maxsize), 2) + #self.assertRaises(ValueError, a.index, 0, 4*sys.maxsize,-4*sys.maxsize) + 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/Modules/arraymodule.c b/Modules/arraymodule.c index 4f778a2dea..7bef363e10 100644 --- a/Modules/arraymodule.c +++ b/Modules/arraymodule.c @@ -1115,18 +1115,42 @@ array_array_count(arrayobject *self, PyObject *v) array.array.index v: object + start: Py_ssize_t = 0 + end: Py_ssize_t = 9223372036854775807 / Return index of first occurrence of v in the array. + +Raises 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]*/ +array_array_index_impl(arrayobject *self, PyObject *v, Py_ssize_t start, + Py_ssize_t end) +/*[clinic end generated code: output=82bf305cb8660fe2 input=da42426e19da3877]*/ { - Py_ssize_t i; + Py_ssize_t i, len; - for (i = 0; i < Py_SIZE(self); i++) { + len = Py_SIZE(self); + + // interpret negative indicies as indexing from the end, in the + // usual way. TODO: Clean up - can this be written better/faster? + if (start < 0) { + start += len; + + // wrap to the start if we are still negative + if (start < 0) { + start = 0; + } + } + + if (end < 0) { + end += len; + } else if (end > len) { + end = len; + } + + for (i = start; i < end; i++) { PyObject *selfi; int cmp; diff --git a/Modules/clinic/arraymodule.c.h b/Modules/clinic/arraymodule.c.h index 2b611951ea..5aa59bf000 100644 --- a/Modules/clinic/arraymodule.c.h +++ b/Modules/clinic/arraymodule.c.h @@ -39,13 +39,37 @@ 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, end=9223372036854775807, /)\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" +"Raises 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)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 end); + +static PyObject * +array_array_index(arrayobject *self, PyObject **args, Py_ssize_t nargs) +{ + PyObject *return_value = NULL; + PyObject *v; + Py_ssize_t start = 0; + Py_ssize_t end = 9223372036854775807; + + if (!_PyArg_ParseStack(args, nargs, "O|nn:index", + &v, &start, &end)) { + goto exit; + } + return_value = array_array_index_impl(self, v, start, end); + +exit: + return return_value; +} PyDoc_STRVAR(array_array_remove__doc__, "remove($self, v, /)\n" @@ -505,4 +529,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=c7dfe61312b236a9 input=a9049054013a1b77]*/ +/*[clinic end generated code: output=a24b7adcee667273 input=a9049054013a1b77]*/