implemented count_iteration

Author: flexatone

performance/__main__.py

@@ -26,6 +26,7 @@
 from performance.reference.util import get_new_indexers_and_screen_ak
 from performance.reference.util import get_new_indexers_and_screen_ref
 from performance.reference.util import split_after_count as split_after_count_ref
+from performance.reference.util import count_iteration as count_iteration_ref
 
 from performance.reference.array_go import ArrayGO as ArrayGOREF
 
@@ -43,6 +44,7 @@
 from arraykit import delimited_to_arrays as delimited_to_arrays_ak
 from arraykit import isna_element as isna_element_ak
 from arraykit import split_after_count as split_after_count_ak
+from arraykit import count_iteration as count_iteration_ak
 
 from arraykit import ArrayGO as ArrayGOAK
 
@@ -716,6 +718,24 @@ class SplitAfterCountREF(SplitAfterCount):
     entry = staticmethod(split_after_count_ref)
 
 
+#-------------------------------------------------------------------------------
+class CountIterations(Perf):
+    NUMBER = 10_000
+
+    def __init__(self):
+        self.strio = io.StringIO('\n'.join(['abcd'] * 10_000))
+
+    def main(self):
+        post = self.entry(self.strio)
+        self.strio.seek(0)
+
+class CountIterationsAK(CountIterations):
+    entry = staticmethod(count_iteration_ak)
+
+class CountIterationsREF(CountIterations):
+    entry = staticmethod(count_iteration_ref)
+
+
 #-------------------------------------------------------------------------------
 
 def get_arg_parser():

performance/reference/util.py

@@ -254,4 +254,12 @@ def split_after_count(string: str, delimiter: str, count: int):
     *left, right = string.split(delimiter, maxsplit=count)
     return ','.join(left), right
 
+def count_iteration(iterable: tp.Iterable):
+    count = 0
+    for i in iterable:
+        count += 1
+    return count
+
+
+
 

src/__init__.py

@@ -20,3 +20,4 @@
 from ._arraykit import iterable_str_to_array_1d as iterable_str_to_array_1d
 from ._arraykit import get_new_indexers_and_screen as get_new_indexers_and_screen
 from ._arraykit import split_after_count as split_after_count
+from ._arraykit import count_iteration as count_iteration

src/__init__.pyi

@@ -52,6 +52,8 @@ def split_after_count(
         count: int,
         ) -> tp.Tuple[str, str]: ...
 
+def count_iteration(__iterable: tp.Iterable) -> int: ...
+
 def immutable_filter(__array: np.ndarray) -> np.ndarray: ...
 def mloc(__array: np.ndarray) -> int: ...
 def name_filter(__name: tp.Hashable) -> tp.Hashable: ...

src/_arraykit.c

@@ -1442,10 +1442,10 @@ AK_CPL_FromIterable(PyObject* iterable, bool type_parse, Py_UCS4 tsep, Py_UCS4 d
         }
         Py_DECREF(field);
     }
+    Py_DECREF(iter);
     if (PyErr_Occurred()) {
         return NULL;
     }
-    Py_DECREF(iter);
     return cpl;
 }
 
@@ -3059,6 +3059,30 @@ split_after_count(PyObject *Py_UNUSED(m), PyObject *args)
 }
 
 
+
+static PyObject *
+count_iteration(PyObject *Py_UNUSED(m), PyObject *iterable)
+{
+    PyObject *iter = PyObject_GetIter(iterable);
+    if (iter == NULL) return NULL;
+
+    int count = 0;
+    PyObject *v;
+
+    while ((v = PyIter_Next(iter))) {
+        count++;
+        Py_DECREF(v);
+    }
+    Py_DECREF(iter);
+    if (PyErr_Occurred()) {
+        return NULL;
+    }
+    PyObject* result = PyLong_FromLong(count);
+    if (result == NULL) return NULL;
+    return result;
+}
+
+
 //------------------------------------------------------------------------------
 
 // Return the integer version of the pointer to underlying data-buffer of array.
@@ -3922,6 +3946,7 @@ static PyMethodDef arraykit_methods[] =  {
             METH_VARARGS | METH_KEYWORDS,
             NULL},
     {"split_after_count", split_after_count, METH_VARARGS, NULL},
+    {"count_iteration", count_iteration, METH_O, NULL},
     {"isna_element", isna_element, METH_O, NULL},
     {"dtype_from_element", dtype_from_element, METH_O, NULL},
     {"get_new_indexers_and_screen",

test/test_util.py

@@ -3,7 +3,7 @@
 import datetime
 import unittest
 import warnings
-
+from io import StringIO
 import numpy as np  # type: ignore
 
 from arraykit import resolve_dtype
@@ -18,6 +18,7 @@
 from arraykit import isna_element
 from arraykit import dtype_from_element
 from arraykit import split_after_count
+from arraykit import count_iteration
 
 from performance.reference.util import get_new_indexers_and_screen_ak as get_new_indexers_and_screen_full
 from arraykit import get_new_indexers_and_screen
@@ -484,5 +485,18 @@ def test_split_after_count_h(self) -> None:
         self.assertEqual(post[0], 'a,b,c,d,e')
         self.assertEqual(post[1], '')
 
+
+    #---------------------------------------------------------------------------
+    def test_count_iteration_a(self) -> None:
+        post = count_iteration(('a', 'b', 'c', 'd'))
+        self.assertEqual(post, 4)
+
+    def test_count_iteration_b(self) -> None:
+        s1 = StringIO(',1,a,b\n-,1,43,54\nX,2,1,3\nY,1,8,10\n-,2,6,20')
+        post = count_iteration(s1)
+        self.assertEqual(post, 5)
+
+
+
 if __name__ == '__main__':
     unittest.main()