diff --git a/src/__init__.py b/src/__init__.py index 9c8bfb8d..1a0648c2 100644 --- a/src/__init__.py +++ b/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 - diff --git a/src/__init__.pyi b/src/__init__.pyi index 888bb58c..1484f9c0 100644 --- a/src/__init__.pyi +++ b/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, ...]]: ... \ No newline at end of file diff --git a/src/_arraykit.c b/src/_arraykit.c index 49b30298..7ce6eb20 100644 --- a/src/_arraykit.c +++ b/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; } - diff --git a/src/methods.c b/src/methods.c index d98f75cd..d2a2543a 100644 --- a/src/methods.c +++ b/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) { diff --git a/src/methods.h b/src/methods.h index 1d33a558..5d2a80c4 100644 --- a/src/methods.h +++ b/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 * diff --git a/test/test_util.py b/test/test_util.py index 7cabbb32..4a58cb61 100644 --- a/test/test_util.py +++ b/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()