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

@ -4149,7 +4149,13 @@ 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, /)
.. method:: count(value, /)
Count the number of occurrences of *value*.
.. versionadded:: next
.. 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*).

View file

@ -90,6 +90,22 @@ class AbstractMemoryTests:
m = self._view(b)
self.assertEqual(list(m), [m[i] for i in range(len(m))])
def test_count(self):
for tp in self._types:
b = tp(self._source)
m = self._view(b)
l = m.tolist()
for ch in list(m):
self.assertEqual(m.count(ch), l.count(ch))
b = tp((b'a' * 5) + (b'c' * 3))
m = self._view(b) # may be sliced
l = m.tolist()
with self.subTest('count', buffer=b):
self.assertEqual(m.count(ord('a')), l.count(ord('a')))
self.assertEqual(m.count(ord('b')), l.count(ord('b')))
self.assertEqual(m.count(ord('c')), l.count(ord('c')))
def test_setitem_readonly(self):
if not self.ro_type:
self.skipTest("no read-only type to test")
@ -464,6 +480,18 @@ class BaseMemoryviewTests:
def _check_contents(self, tp, obj, contents):
self.assertEqual(obj, tp(contents))
def test_count(self):
super().test_count()
for tp in self._types:
b = tp((b'a' * 5) + (b'c' * 3))
m = self._view(b) # should not be sliced
self.assertEqual(len(b), len(m))
with self.subTest('count', buffer=b):
self.assertEqual(m.count(ord('a')), 5)
self.assertEqual(m.count(ord('b')), 0)
self.assertEqual(m.count(ord('c')), 3)
class BaseMemorySliceTests:
source_bytes = b"XabcdefY"

View file

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

View file

@ -419,6 +419,15 @@ exit:
return return_value;
}
PyDoc_STRVAR(memoryview_count__doc__,
"count($self, value, /)\n"
"--\n"
"\n"
"Count the number of occurrences of a value.");
#define MEMORYVIEW_COUNT_METHODDEF \
{"count", (PyCFunction)memoryview_count, METH_O, memoryview_count__doc__},
PyDoc_STRVAR(memoryview_index__doc__,
"index($self, value, start=0, stop=sys.maxsize, /)\n"
"--\n"
@ -464,4 +473,4 @@ skip_optional:
exit:
return return_value;
}
/*[clinic end generated code: output=2742d371dba7314f input=a9049054013a1b77]*/
/*[clinic end generated code: output=132893ef5f67ad73 input=a9049054013a1b77]*/

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},