additional testing

Author: flexatone

doc/articles/block_index.py

@@ -75,6 +75,7 @@ def __init__(self, arrays: tp.Iterable[np.ndarray]):
 
         self.selector_int_array = np.arange(0, len(self.bi), 2)
         self.selector_int_list = list(range(0, len(self.bi), 2))
+        self.selector_bool_array = (np.arange(len(self.bi)) % 2) == 0
 
 #-------------------------------------------------------------------------------
 class BlockIndexLoad(ArrayProcessor):
@@ -191,6 +192,23 @@ def __call__(self):
 
 
 
+class BlockIndexIterBoolArray(ArrayProcessor):
+    NAME = 'BlockIndex: iter by bool array'
+    SORT = 7
+
+    def __call__(self):
+        _ = list(self.bi.iter_select(self.selector_bool_array))
+
+class TupleIndexIterBoolArray(ArrayProcessor):
+    NAME = 'TupleIndex: iter by bool array'
+    SORT = 17
+
+    def __call__(self):
+        ti = self.ti
+        _ = [ti[i] for i in self.selector_bool_array if i]
+
+
+
 
 #-------------------------------------------------------------------------------
 NUMBER = 2
@@ -360,6 +378,9 @@ def get_versions() -> str:
     TupleIndexIterIntArray,
     BlockIndexIterIntList,
     TupleIndexIterIntList,
+    BlockIndexIterBoolArray,
+    TupleIndexIterBoolArray,
+
     )
 
 CLS_FF = (

src/_arraykit.c

@@ -4146,7 +4146,7 @@ typedef struct BIIterObject {
     PyObject_VAR_HEAD
     BlockIndexObject *bi;
     int8_t reversed;
-    Py_ssize_t index; // current index state, mutated in-place
+    Py_ssize_t pos; // current index state, mutated in-place
 } BIIterObject;
 
 static PyObject *
@@ -4158,7 +4158,7 @@ BIIter_new(BlockIndexObject *bi, int8_t reversed) {
     Py_INCREF(bi);
     bii->bi = bi;
     bii->reversed = reversed;
-    bii->index = 0;
+    bii->pos = 0;
     return (PyObject *)bii;
 }
 
@@ -4178,13 +4178,13 @@ static PyObject *
 BIIter_iternext(BIIterObject *self) {
     Py_ssize_t i;
     if (self->reversed) {
-        i = self->bi->bir_count - ++self->index;
+        i = self->bi->bir_count - ++self->pos;
         if (i < 0) {
             return NULL;
         }
     }
     else {
-        i = self->index++;
+        i = self->pos++;
     }
     if (self->bi->bir_count <= i) {
         return NULL;
@@ -4199,8 +4199,8 @@ BIIter_reversed(BIIterObject *self) {
 
 static PyObject *
 BIIter_length_hint(BIIterObject *self) {
-    // this works for reversed as we use self-> index to subtract from length
-    Py_ssize_t len = Py_MAX(0, self->bi->bir_count - self->index);
+    // this works for reversed as we use self->pos to subtract from length
+    Py_ssize_t len = Py_MAX(0, self->bi->bir_count - self->pos);
     return PyLong_FromSsize_t(len);
 }
 
@@ -4564,13 +4564,15 @@ BIIterSelector_new(BlockIndexObject *bi,
             PyErr_SetString(PyExc_TypeError, "Slices cannot be used as selectors for this type of iterator");
             return NULL;
         }
-        if (PySlice_GetIndicesEx(selector, bi->bir_count, &pos, &stop, &step, &len)) {
+        if (PySlice_Unpack(selector, &pos, &stop, &step)) {
             return NULL;
         }
+        len = PySlice_AdjustIndices(bi->bir_count, &pos, &stop, step);
         if (reversed) {
             pos += (step * (len - 1));
             step *= -1;
         }
+        // AK_DEBUG_MSG_OBJ("resolved slice", Py_BuildValue("nnnn", pos, stop, step, len));
     }
     else if (PyList_CheckExact(selector)) {
         if (kind == BIIS_UNKNOWN) {

test/test_block_index.py

@@ -409,8 +409,40 @@ def test_block_index_iter_select_slice_c(self) -> None:
             [(2, 1), (2, 0), (1, 0), (0, 1)]
             )
 
+
+    def test_block_index_iter_select_slice_d(self) -> None:
+        bi1 = BlockIndex()
+        bi1.register(np.arange(6).reshape(2,3))
+        bi1.register(np.arange(2))
+
+        self.assertEqual(list(bi1.iter_select(slice(None))),
+            [(0, 0), (0, 1), (0, 2), (1, 0)]
+            )
+        self.assertEqual(list(bi1.iter_select(slice(20, 24))),
+            []
+            )
+        self.assertEqual(list(bi1.iter_select(slice(0, 100, 10))),
+            [(0, 0)]
+            )
+        self.assertEqual(list(bi1.iter_select(slice(0, 100, 3))),
+            [(0, 0), (1, 0)]
+            )
+
+    def test_block_index_iter_select_slice_e(self) -> None:
+        bi1 = BlockIndex()
+        bi1.register(np.arange(12).reshape(2,6))
+        bi1.register(np.arange(12).reshape(2,6))
+
+        self.assertEqual(list(bi1.iter_select(slice(11, None, -3))),
+            [(1, 5), (1, 2), (0, 5), (0, 2)]
+            )
+        self.assertEqual(list(bi1.iter_select(slice(11, None, -4))),
+            [(1, 5), (1, 1), (0, 3)]
+            )
+
+
     #---------------------------------------------------------------------------
-    def test_block_index_iter_select_slice_a(self) -> None:
+    def test_block_index_iter_select_boolean_a(self) -> None:
         bi1 = BlockIndex()
         bi1.register(np.arange(4).reshape(2,2))
         bi1.register(np.arange(2))
@@ -428,10 +460,25 @@ def test_block_index_iter_select_slice_a(self) -> None:
                 [(0, 0), (2, 4)]
                 )
 
-    def test_block_index_iter_select_slice_b(self) -> None:
+    def test_block_index_iter_select_boolean_b(self) -> None:
         bi1 = BlockIndex()
         bi1.register(np.arange(4).reshape(2,2))
         bi1.register(np.arange(2))
 
         with self.assertRaises(TypeError):
-            bi1.iter_select(np.array([False, True]))
+            bi1.iter_select(np.array([False, True]))
+
+        with self.assertRaises(TypeError):
+            bi1.iter_select(np.full(20, True))
+
+
+    def test_block_index_iter_select_boolean_c(self) -> None:
+        bi1 = BlockIndex()
+        bi1.register(np.arange(4).reshape(2,2))
+        bi1.register(np.arange(2))
+        self.assertEqual(list(bi1.iter_select(np.full(len(bi1), False))),
+                []
+                )
+        self.assertEqual(list(bi1.iter_select(np.full(len(bi1), True))),
+                [(0, 0), (0, 1), (1, 0)]
+                )