gh-125420: implement Sequence.index API on memoryview objects (#125446)

This commit is contained in:
Bénédikt Tran 2024-12-10 03:48:38 +01:00 committed by GitHub
parent 3b18af964d
commit 58c753827a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 171 additions and 1 deletions

View File

@ -4149,6 +4149,15 @@ copying.
.. versionchanged:: 3.5
The source format is no longer restricted when casting to a byte view.
.. method:: index(value, start=0, stop=sys.maxsize, /)
Return the index of the first occurrence of *value* (at or after
index *start* and before index *stop*).
Raises a :exc:`ValueError` if *value* cannot be found.
.. versionadded:: next
There are also several readonly attributes available:
.. attribute:: obj

View File

@ -15,6 +15,7 @@ import copy
import pickle
import struct
from itertools import product
from test.support import import_helper
@ -58,6 +59,31 @@ class AbstractMemoryTests:
for tp in self._types:
self.check_getitem_with_type(tp)
def test_index(self):
for tp in self._types:
b = tp(self._source)
m = self._view(b) # may be a sub-view
l = m.tolist()
k = 2 * len(self._source)
for chi in self._source:
if chi in l:
self.assertEqual(m.index(chi), l.index(chi))
else:
self.assertRaises(ValueError, m.index, chi)
for start, stop in product(range(-k, k), range(-k, k)):
index = -1
try:
index = l.index(chi, start, stop)
except ValueError:
pass
if index == -1:
self.assertRaises(ValueError, m.index, chi, start, stop)
else:
self.assertEqual(m.index(chi, start, stop), index)
def test_iter(self):
for tp in self._types:
b = tp(self._source)

View File

@ -0,0 +1,2 @@
Add :meth:`memoryview.index` to :class:`memoryview` objects. Patch by
Bénédikt Tran.

View File

@ -418,4 +418,50 @@ skip_optional_pos:
exit:
return return_value;
}
/*[clinic end generated code: output=0a93f08110630633 input=a9049054013a1b77]*/
PyDoc_STRVAR(memoryview_index__doc__,
"index($self, value, start=0, stop=sys.maxsize, /)\n"
"--\n"
"\n"
"Return the index of the first occurrence of a value.\n"
"\n"
"Raises ValueError if the value is not present.");
#define MEMORYVIEW_INDEX_METHODDEF \
{"index", _PyCFunction_CAST(memoryview_index), METH_FASTCALL, memoryview_index__doc__},
static PyObject *
memoryview_index_impl(PyMemoryViewObject *self, PyObject *value,
Py_ssize_t start, Py_ssize_t stop);
static PyObject *
memoryview_index(PyMemoryViewObject *self, PyObject *const *args, Py_ssize_t nargs)
{
PyObject *return_value = NULL;
PyObject *value;
Py_ssize_t start = 0;
Py_ssize_t stop = PY_SSIZE_T_MAX;
if (!_PyArg_CheckPositional("index", nargs, 1, 3)) {
goto exit;
}
value = 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 = memoryview_index_impl(self, value, start, stop);
exit:
return return_value;
}
/*[clinic end generated code: output=2742d371dba7314f input=a9049054013a1b77]*/

View File

@ -2748,6 +2748,92 @@ static PySequenceMethods memory_as_sequence = {
};
/**************************************************************************/
/* Lookup */
/**************************************************************************/
/*[clinic input]
memoryview.index
value: object
start: slice_index(accept={int}) = 0
stop: slice_index(accept={int}, c_default="PY_SSIZE_T_MAX") = sys.maxsize
/
Return the index of the first occurrence of a value.
Raises ValueError if the value is not present.
[clinic start generated code]*/
static PyObject *
memoryview_index_impl(PyMemoryViewObject *self, PyObject *value,
Py_ssize_t start, Py_ssize_t stop)
/*[clinic end generated code: output=e0185e3819e549df input=0697a0165bf90b5a]*/
{
const Py_buffer *view = &self->view;
CHECK_RELEASED(self);
if (view->ndim == 0) {
PyErr_SetString(PyExc_TypeError, "invalid lookup on 0-dim memory");
return NULL;
}
if (view->ndim == 1) {
Py_ssize_t n = view->shape[0];
if (start < 0) {
start = Py_MAX(start + n, 0);
}
if (stop < 0) {
stop = Py_MAX(stop + n, 0);
}
stop = Py_MIN(stop, n);
assert(stop >= 0);
assert(stop <= n);
start = Py_MIN(start, stop);
assert(0 <= start);
assert(start <= stop);
PyObject *obj = _PyObject_CAST(self);
for (Py_ssize_t index = start; index < stop; index++) {
// Note: while memoryviews can be mutated during iterations
// when calling the == operator, their shape cannot. As such,
// it is safe to assume that the index remains valid for the
// entire loop.
assert(index < n);
PyObject *item = memory_item(obj, index);
if (item == NULL) {
return NULL;
}
if (item == value) {
Py_DECREF(item);
return PyLong_FromSsize_t(index);
}
int contained = PyObject_RichCompareBool(item, value, Py_EQ);
Py_DECREF(item);
if (contained > 0) { // more likely than 'contained < 0'
return PyLong_FromSsize_t(index);
}
else if (contained < 0) {
return NULL;
}
}
PyErr_SetString(PyExc_ValueError, "memoryview.index(x): x not found");
return NULL;
}
PyErr_SetString(PyExc_NotImplementedError,
"multi-dimensional lookup is not implemented");
return NULL;
}
/**************************************************************************/
/* Comparisons */
/**************************************************************************/
@ -3284,6 +3370,7 @@ static PyMethodDef memory_methods[] = {
MEMORYVIEW_CAST_METHODDEF
MEMORYVIEW_TOREADONLY_METHODDEF
MEMORYVIEW__FROM_FLAGS_METHODDEF
MEMORYVIEW_INDEX_METHODDEF
{"__enter__", memory_enter, METH_NOARGS, NULL},
{"__exit__", memory_exit, METH_VARARGS, memory_exit_doc},
{"__class_getitem__", Py_GenericAlias, METH_O|METH_CLASS, PyDoc_STR("See PEP 585")},