updated register implementation, added tests

Author: flexatone

doc/articles/block_index.py

@@ -9,7 +9,7 @@
 import pickle
 
 from arraykit import BlockIndex
-# from arraykit import ErrorInitBlocks
+# from arraykit import ErrorInitTypeBlocks
 from arraykit import shape_filter
 from arraykit import resolve_dtype
 

src/__init__.py

@@ -5,7 +5,7 @@
 from ._arraykit import __version__
 from ._arraykit import ArrayGO as ArrayGO
 from ._arraykit import BlockIndex as BlockIndex
-from ._arraykit import ErrorInitBlocks as ErrorInitBlocks
+from ._arraykit import ErrorInitTypeBlocks as ErrorInitTypeBlocks
 
 from ._arraykit import immutable_filter as immutable_filter
 from ._arraykit import mloc as mloc

src/__init__.pyi

@@ -6,7 +6,7 @@ _T = tp.TypeVar('_T')
 
 __version__: str
 
-class ErrorInitBlocks:
+class ErrorInitTypeBlocks:
     def __init__(self, *args: tp.Any, **kwargs: tp.Any) -> None: ...
     def with_traceback(self, tb: Exception) -> Exception: ...
     def __setstate__(self) -> None: ...
@@ -27,10 +27,12 @@ class ArrayGO:
 
 class BlockIndex:
     shape: tp.Tuple[int, int]
-    dtype: tp.Optional[np.dtype]
+    dtype: np.dtype
+    rows: int
+    columns: int
 
     def __init__() -> None: ...
-    def register(self, __value: object) -> None: ...
+    def register(self, __value: np.ndarray) -> bool: ...
     def to_list(self,) -> tp.List[int]: ...
     def to_bytes(self,) -> bytes: ...
     def copy(self,) -> 'BlockIndex': ...

src/_arraykit.c

@@ -4114,7 +4114,7 @@ get_new_indexers_and_screen(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kw
 //------------------------------------------------------------------------------
 
 static PyTypeObject BlockIndexType;
-static PyObject *ErrorInitBlocks;
+static PyObject *ErrorInitTypeBlocks;
 
 // NOTE: we use platform size types here, which are appropriate for the values, but might pose issues if trying to pass pickles between 32 and 64 bit machines.
 typedef struct BlockIndexRecord {
@@ -4777,18 +4777,18 @@ BlockIndex_repr(BlockIndexObject *self) {
             dt);
 }
 
-// Returns NULL on error, None otherwise. This checks and raises on non-array inputs, dimensions other than 1 or 2.
+// Returns NULL on error, True if the block should be reatained, False if the block has zero columns and should not be retained. This checks and raises on non-array inputs, dimensions other than 1 or 2, and mis-aligned columns.
 static PyObject *
 BlockIndex_register(BlockIndexObject *self, PyObject *value) {
     if (!PyArray_Check(value)) {
-        PyErr_Format(ErrorInitBlocks, "Found non-array block: %R", value);
+        PyErr_Format(ErrorInitTypeBlocks, "Found non-array block: %R", value);
         return NULL;
     }
     PyArrayObject *a = (PyArrayObject *)value;
     int ndim = PyArray_NDIM(a);
 
     if (ndim < 1 || ndim > 2) {
-        PyErr_Format(ErrorInitBlocks, "Array block has invalid dimensions: %i", ndim);
+        PyErr_Format(ErrorInitTypeBlocks, "Array block has invalid dimensions: %i", ndim);
         return NULL;
     }
     Py_ssize_t increment = ndim == 1 ? 1 : PyArray_DIM(a, 1);
@@ -4799,13 +4799,17 @@ BlockIndex_register(BlockIndexObject *self, PyObject *value) {
         self->row_count = alignment;
     }
     else if (self->row_count != alignment) {
-        PyErr_Format(ErrorInitBlocks,
+        PyErr_Format(ErrorInitTypeBlocks,
                 "Array block has unaligned row count: found %i, expected %i",
                 alignment,
                 self->row_count);
         return NULL;
     }
 
+    if (increment == 0) {
+        Py_RETURN_FALSE;
+    }
+
     PyArray_Descr* dt = PyArray_DESCR(a); // borrowed ref
     if (self->dtype == NULL) {
         Py_INCREF((PyObject*)dt);
@@ -4829,7 +4833,7 @@ BlockIndex_register(BlockIndexObject *self, PyObject *value) {
         self->bir_count++;
     }
     self->block_count++;
-    Py_RETURN_NONE;
+    Py_RETURN_TRUE;
 }
 
 
@@ -4942,6 +4946,7 @@ BlockIndex_iter(BlockIndexObject* self) {
 
 static PyObject *
 BlockIndex_shape_getter(BlockIndexObject *self, void* Py_UNUSED(closure)){
+    // NOTE: this could be cached
     return Py_BuildValue("nn", self->row_count, self->bir_count);
 }
 
@@ -5395,12 +5400,12 @@ PyInit__arraykit(void)
 {
     import_array();
 
-    ErrorInitBlocks = PyErr_NewExceptionWithDoc(
-            "arraykit.ErrorInitBlocks",
+    ErrorInitTypeBlocks = PyErr_NewExceptionWithDoc(
+            "arraykit.ErrorInitTypeBlocks",
             "RuntimeError error in block initialization.",
             PyExc_RuntimeError,
             NULL);
-    if (ErrorInitBlocks == NULL) {
+    if (ErrorInitTypeBlocks == NULL) {
         return NULL;
     }
 
@@ -5426,7 +5431,7 @@ PyInit__arraykit(void)
         PyModule_AddObject(m, "BlockIndex", (PyObject *) &BlockIndexType) ||
         PyModule_AddObject(m, "ArrayGO", (PyObject *) &ArrayGOType) ||
         PyModule_AddObject(m, "deepcopy", deepcopy) ||
-        PyModule_AddObject(m, "ErrorInitBlocks", ErrorInitBlocks)
+        PyModule_AddObject(m, "ErrorInitTypeBlocks", ErrorInitTypeBlocks)
     ){
         Py_DECREF(deepcopy);
         Py_XDECREF(m);

test/test_block_index.py

@@ -7,7 +7,7 @@
 import numpy as np
 
 from arraykit import BlockIndex
-from arraykit import ErrorInitBlocks
+from arraykit import ErrorInitTypeBlocks
 
 
 class TestUnit(unittest.TestCase):
@@ -53,19 +53,19 @@ def test_block_index_init_d(self) -> None:
 
     def test_block_index_register_a(self) -> None:
         bi1 = BlockIndex()
-        with self.assertRaises(ErrorInitBlocks):
+        with self.assertRaises(ErrorInitTypeBlocks):
             bi1.register('foo')
 
-        with self.assertRaises(ErrorInitBlocks):
+        with self.assertRaises(ErrorInitTypeBlocks):
             bi1.register(3.5)
 
     def test_block_index_register_b(self) -> None:
 
         bi1 = BlockIndex()
-        with self.assertRaises(ErrorInitBlocks):
+        with self.assertRaises(ErrorInitTypeBlocks):
             bi1.register(np.array(0))
 
-        with self.assertRaises(ErrorInitBlocks):
+        with self.assertRaises(ErrorInitTypeBlocks):
             bi1.register(np.arange(12).reshape(2,3,2))
 
 
@@ -96,7 +96,7 @@ def test_block_index_register_d(self) -> None:
     def test_block_index_register_e(self) -> None:
         bi1 = BlockIndex()
         bi1.register(np.arange(2))
-        with self.assertRaises(ErrorInitBlocks):
+        with self.assertRaises(ErrorInitTypeBlocks):
             bi1.register(np.arange(12).reshape(3,4))
 
 
@@ -108,6 +108,39 @@ def test_block_index_register_f(self) -> None:
         self.assertEqual(bi1.columns, 10_000)
 
 
+    def test_block_index_register_g(self) -> None:
+        bi1 = BlockIndex()
+        a1 = np.array(()).reshape(4, 0)
+        self.assertFalse(bi1.register(a1))
+        self.assertEqual(bi1.shape, (4, 0))
+        # as not dtype has been registered, we will get default float
+        self.assertEqual(bi1.dtype, np.dtype(float))
+
+        a2 = np.arange(8).reshape(4, 2).astype(bool)
+        self.assertTrue(bi1.register(a2))
+        self.assertEqual(bi1.shape, (4, 2))
+        self.assertEqual(bi1.dtype, np.dtype(bool))
+
+
+    def test_block_index_register_h(self) -> None:
+        bi1 = BlockIndex()
+        a1 = np.array(()).reshape(0, 4).astype(bool)
+        self.assertTrue(bi1.register(a1))
+        self.assertEqual(bi1.shape, (0, 4))
+        self.assertEqual(bi1.dtype, np.dtype(bool))
+
+        a2 = np.array(()).reshape(0, 0).astype(float)
+        self.assertFalse(bi1.register(a2))
+        self.assertEqual(bi1.shape, (0, 4))
+        # dtype is still bool
+        self.assertEqual(bi1.dtype, np.dtype(bool))
+
+        a3 = np.array(()).reshape(0, 3).astype(int)
+        self.assertTrue(bi1.register(a3))
+        self.assertEqual(bi1.shape, (0, 7))
+        self.assertEqual(bi1.dtype, np.dtype(object))
+
+
     #---------------------------------------------------------------------------
 
     def test_block_index_to_bytes_a(self) -> None:

test/test_pyi.py

@@ -30,7 +30,7 @@ def from_module(cls, module):
                 continue
             obj = getattr(module, name)
             if isinstance(obj, type): # a class
-                if name == ak.ErrorInitBlocks.__name__:
+                if name == ak.ErrorInitTypeBlocks.__name__:
                     # skip as there is Python version variability
                     continue
                 classes[name] = []