Merge pull request #209 from static-frame/copilot/slice-to-unit-function

Add `slice_to_unit()` C API for single-element slice detection

Author: flexatone

src/__init__.py

@@ -28,6 +28,7 @@
 from ._arraykit import first_true_1d as first_true_1d
 from ._arraykit import first_true_2d as first_true_2d
 from ._arraykit import slice_to_ascending_slice as slice_to_ascending_slice
+from ._arraykit import slice_to_unit as slice_to_unit
 from ._arraykit import array_to_tuple_array as array_to_tuple_array
 from ._arraykit import array_to_tuple_iter as array_to_tuple_iter
 from ._arraykit import nonzero_1d as nonzero_1d
@@ -37,4 +38,3 @@
 from ._arraykit import AutoMap as AutoMap
 from ._arraykit import FrozenAutoMap as FrozenAutoMap
 from ._arraykit import NonUniqueError as NonUniqueError
-

src/__init__.pyi

@@ -207,5 +207,6 @@ def is_objectable_dt64(__array: np.ndarray, /) -> bool: ...
 def is_objectable(__array: np.ndarray, /) -> bool: ...
 def astype_array(__array: np.ndarray, __dtype: np.dtype | None, /) -> np.ndarray: ...
 def slice_to_ascending_slice(__slice: slice, __size: int) -> slice: ...
+def slice_to_unit(__slice: slice, /) -> int: ...
 def array_to_tuple_array(__array: np.ndarray) -> np.ndarray: ...
 def array_to_tuple_iter(__array: np.ndarray) -> tp.Iterator[tp.Tuple[tp.Any, ...]]: ...

src/_arraykit.c

@@ -22,6 +22,7 @@ static PyMethodDef arraykit_methods[] =  {
     {"column_1d_filter", column_1d_filter, METH_O, NULL},
     {"row_1d_filter", row_1d_filter, METH_O, NULL},
     {"slice_to_ascending_slice", slice_to_ascending_slice, METH_VARARGS, NULL},
+    {"slice_to_unit", slice_to_unit, METH_O, NULL},
     {"array_deepcopy",
             (PyCFunction)array_deepcopy,
             METH_VARARGS | METH_KEYWORDS,
@@ -157,4 +158,3 @@ PyInit__arraykit(void)
 #endif
     return m;
 }
-

src/methods.c

@@ -60,10 +60,55 @@ slice_to_ascending_slice(PyObject *Py_UNUSED(m), PyObject *args) {
             &PyLong_Type, &size)) {
         return NULL;
     }
-    // will delegate NULL on eroror
+    // will delegate NULL on error
     return AK_slice_to_ascending_slice(slice, PyLong_AsSsize_t(size));
 }
 
+PyObject *
+slice_to_unit(PyObject *Py_UNUSED(m), PyObject *a)
+{
+    if (!PySlice_Check(a)) {
+        return PyErr_Format(PyExc_TypeError,
+                "Expected a slice, not %s",
+                Py_TYPE(a)->tp_name);
+    }
+    PyObject* py_start = ((PySliceObject*)a)->start;
+    PyObject* py_stop = ((PySliceObject*)a)->stop;
+    PyObject* py_step = ((PySliceObject*)a)->step;
+
+    if (py_stop == Py_None) {
+        return PyLong_FromLong(-1);
+    }
+
+    Py_ssize_t step = 1;
+    if (py_step != Py_None) {
+        step = PyLong_AsSsize_t(py_step);
+        if (step == -1 && PyErr_Occurred()) {
+            return NULL;
+        }
+    }
+    if (step != 1) {
+        return PyLong_FromLong(-1);
+    }
+
+    Py_ssize_t start = 0;
+    if (py_start != Py_None) {
+        start = PyLong_AsSsize_t(py_start);
+        if (start == -1 && PyErr_Occurred()) {
+            return NULL;
+        }
+    }
+    Py_ssize_t stop = PyLong_AsSsize_t(py_stop);
+    if (stop == -1 && PyErr_Occurred()) {
+        return NULL;
+    }
+
+    if (start < 0 || stop < 0 || stop - start != 1) {
+        return PyLong_FromLong(-1);
+    }
+    return PyLong_FromSsize_t(start);
+}
+
 PyObject *
 column_2d_filter(PyObject *Py_UNUSED(m), PyObject *a)
 {

src/methods.h

@@ -14,6 +14,10 @@ row_1d_filter(PyObject *Py_UNUSED(m), PyObject *a);
 PyObject *
 slice_to_ascending_slice(PyObject *Py_UNUSED(m), PyObject *args);
 
+// Return an integer when a slice is exactly a single positive-position unit, else -1.
+PyObject *
+slice_to_unit(PyObject *Py_UNUSED(m), PyObject *a);
+
 // Reshape if necessary a flat ndim 1 array into a 2D array with one columns and rows of length.
 // related example: https://github.com/RhysU/ar/blob/master/ar-python.cpp
 PyObject *

test/test_util.py

@@ -23,6 +23,7 @@
 from arraykit import first_true_1d
 from arraykit import first_true_2d
 from arraykit import slice_to_ascending_slice
+from arraykit import slice_to_unit
 from arraykit import array_to_tuple_array
 from arraykit import array_to_tuple_iter
 
@@ -953,6 +954,24 @@ def test_slice_to_ascending_slice_i(self) -> None:
             slice(1, 2, None)
             )
 
+    def test_slice_to_unit_a(self) -> None:
+        self.assertEqual(slice_to_unit(slice(3, 4)), 3)
+        self.assertEqual(slice_to_unit(slice(0, 1)), 0)
+        self.assertEqual(slice_to_unit(slice(None, 1)), 0)
+
+    def test_slice_to_unit_b(self) -> None:
+        self.assertEqual(slice_to_unit(slice(0, 2)), -1)
+        self.assertEqual(slice_to_unit(slice(5, 5)), -1)
+        self.assertEqual(slice_to_unit(slice(0, 1, 2)), -1)
+        self.assertEqual(slice_to_unit(slice(0, 1, -1)), -1)
+        self.assertEqual(slice_to_unit(slice(0, None)), -1)
+        self.assertEqual(slice_to_unit(slice(None, 2)), -1)
+        self.assertEqual(slice_to_unit(slice(-1, 0)), -1)
+
+    def test_slice_to_unit_c(self) -> None:
+        with self.assertRaises(TypeError):
+            _ = slice_to_unit(3)
+
 
 if __name__ == '__main__':
     unittest.main()