Merge pull request #44 from peytondmurray/pyarray-arrfunc-nonzero

Add support for `PyArray_NonzeroFunc`

Author: ngoldbaum

stringdtype/stringdtype/src/dtype.c

@@ -168,6 +168,14 @@ stringdtype_getitem(StringDTypeObject *NPY_UNUSED(descr), char **dataptr)
     return res;
 }
 
+// PyArray_NonzeroFunc
+// Unicode strings are nonzero if their length is nonzero.
+npy_bool
+nonzero(void *data, void *NPY_UNUSED(arr))
+{
+    return ((ss *)data)->len != 0;
+}
+
 // Implementation of PyArray_CompareFunc.
 // Compares unicode strings by their code points.
 int
@@ -225,6 +233,7 @@ static PyType_Slot StringDType_Slots[] = {
         {NPY_DT_getitem, &stringdtype_getitem},
         {NPY_DT_ensure_canonical, &stringdtype_ensure_canonical},
         {NPY_DT_PyArray_ArrFuncs_compare, &compare_strings},
+        {NPY_DT_PyArray_ArrFuncs_nonzero, &nonzero},
         {NPY_DT_get_clear_loop, &stringdtype_get_clear_loop},
         {0, NULL}};
 

stringdtype/tests/test_stringdtype.py

@@ -183,6 +183,20 @@ def test_sort(strings):
     np.testing.assert_array_equal(arr, arr_sorted)
 
 
+@pytest.mark.parametrize(
+    "strings",
+    [
+        ["A¢☃€ 😊", " A☃€¢😊", "☃€😊 A¢", "😊☃A¢ €"],
+        ["A¢☃€ 😊", "", " ", " "],
+        ["", "a", "😸", "ááðfáíóåéë"],
+    ],
+)
+def test_nonzero(strings):
+    arr = np.array(strings, dtype=StringDType())
+    is_nonzero = np.array([i for i, item in enumerate(arr) if len(item) != 0])
+    np.testing.assert_array_equal(arr.nonzero()[0], is_nonzero)
+
+
 def test_creation_functions():
     np.testing.assert_array_equal(
         np.zeros(3, dtype=StringDType()), ["", "", ""]