Merge pull request #122 from static-frame/121/block-index-neg-indices

Sequence iters to handle negative numbers

Author: flexatone

src/_arraykit.c

@@ -4183,8 +4183,8 @@ typedef struct BlockIndexObject {
 } BlockIndexObject;
 
 
-// Returns a new reference to tuple. Returns NULL on error.
-static PyObject*
+// Returns a new reference to tuple. Returns NULL on error. Python already wraps negative numbers up to negative length when used in the sequence slot
+static inline PyObject*
 AK_BI_item(BlockIndexObject* self, Py_ssize_t i) {
     if (!((size_t)i < (size_t)self->bir_count)) {
         PyErr_SetString(PyExc_IndexError, "index out of range");
@@ -4194,6 +4194,16 @@ AK_BI_item(BlockIndexObject* self, Py_ssize_t i) {
     return Py_BuildValue("nn", biri->block, biri->column); // maybe NULL
 }
 
+// Returns a new reference to tuple. Returns NULL on error. Supports negative numbers up to negative length.
+static inline PyObject*
+AK_BI_item_wraps(BlockIndexObject* self, Py_ssize_t i)
+{
+    if (i < 0) {
+        i = self->bir_count + i;
+    }
+    return AK_BI_item(self, i);
+}
+
 //------------------------------------------------------------------------------
 // BI Iterator
 static PyTypeObject BIIterType;
@@ -4378,7 +4388,7 @@ BIIterSeq_iternext(BIIterSeqObject *self) {
             return NULL;
         }
     }
-    return AK_BI_item(self->bi, t); // return new ref
+    return AK_BI_item_wraps(self->bi, t); // return new ref
 }
 
 static PyObject *
@@ -4557,7 +4567,6 @@ static PyTypeObject BIIterBooleanType = {
     .tp_name = "arraykit.BlockIndexIteratorBoolean",
 };
 
-
 //------------------------------------------------------------------------------
 
 // NOTE: this constructor returns one of three different PyObject types. We do this to consolidate error reporting and type checks.
@@ -4628,7 +4637,6 @@ BIIterSelector_new(BlockIndexObject *bi,
             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) {
@@ -4694,8 +4702,6 @@ BIIterSelector_new(BlockIndexObject *bi,
     return NULL;
 }
 
-
-
 //------------------------------------------------------------------------------
 
 // Returns 0 on succes, -1 on error.

test/test_block_index.py

@@ -528,4 +528,27 @@ def test_block_index_iter_select_boolean_c(self) -> None:
                 )
         self.assertEqual(list(bi1.iter_select(np.full(len(bi1), True))),
                 [(0, 0), (0, 1), (1, 0)]
-                )
+                )
+
+    #---------------------------------------------------------------------------
+    def test_block_index_iter_select_sequence_a(self) -> None:
+        bi1 = BlockIndex()
+        bi1.register(np.arange(4).reshape(2,2))
+        bi1.register(np.arange(2))
+        bi1.register(np.arange(10).reshape(2,5))
+
+        self.assertEqual(list(bi1.iter_select([0, -1, -2, -8])),
+                [(0, 0), (2, 4), (2, 3), (0, 0)]
+                )
+        self.assertEqual(list(bi1.iter_select(np.array([0, -1, -2, -8]))),
+                [(0, 0), (2, 4), (2, 3), (0, 0)]
+                )
+
+    def test_block_index_iter_select_sequence_b(self) -> None:
+        bi1 = BlockIndex()
+        bi1.register(np.arange(4).reshape(2,2))
+        bi1.register(np.arange(2))
+        bi1.register(np.arange(10).reshape(2,5))
+
+        with self.assertRaises(IndexError):
+            _ = list(bi1.iter_select([-9]))