gh-125420: implement Sequence.count API on memoryview objects (#125443)

This commit is contained in:
Bénédikt Tran 2024-12-10 11:12:33 +01:00 committed by GitHub
parent 050d59bd17
commit 4331832db0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 97 additions and 2 deletions

View file

@ -2748,6 +2748,55 @@ static PySequenceMethods memory_as_sequence = {
};
/****************************************************************************/
/* Counting */
/****************************************************************************/
/*[clinic input]
memoryview.count
value: object
/
Count the number of occurrences of a value.
[clinic start generated code]*/
static PyObject *
memoryview_count(PyMemoryViewObject *self, PyObject *value)
/*[clinic end generated code: output=e2c255a8d54eaa12 input=e3036ce1ed7d1823]*/
{
PyObject *iter = PyObject_GetIter(_PyObject_CAST(self));
if (iter == NULL) {
return NULL;
}
Py_ssize_t count = 0;
PyObject *item = NULL;
while (PyIter_NextItem(iter, &item)) {
if (item == NULL) {
Py_DECREF(iter);
return NULL;
}
if (item == value) {
Py_DECREF(item);
count++; // no overflow since count <= len(mv) <= PY_SSIZE_T_MAX
continue;
}
int contained = PyObject_RichCompareBool(item, value, Py_EQ);
Py_DECREF(item);
if (contained > 0) { // more likely than 'contained < 0'
count++; // no overflow since count <= len(mv) <= PY_SSIZE_T_MAX
}
else if (contained < 0) {
Py_DECREF(iter);
return NULL;
}
}
Py_DECREF(iter);
return PyLong_FromSsize_t(count);
}
/**************************************************************************/
/* Lookup */
/**************************************************************************/
@ -3370,6 +3419,7 @@ static PyMethodDef memory_methods[] = {
MEMORYVIEW_CAST_METHODDEF
MEMORYVIEW_TOREADONLY_METHODDEF
MEMORYVIEW__FROM_FLAGS_METHODDEF
MEMORYVIEW_COUNT_METHODDEF
MEMORYVIEW_INDEX_METHODDEF
{"__enter__", memory_enter, METH_NOARGS, NULL},
{"__exit__", memory_exit, METH_VARARGS, memory_exit_doc},