Merge pull request #46 from ngoldbaum/cast-to-bool

Add a cast from string to bool

Author: ngoldbaum

stringdtype/stringdtype/src/casts.c

@@ -12,6 +12,8 @@ gil_error(PyObject *type, const char *msg)
     PyGILState_Release(gstate);
 }
 
+// string to string
+
 static NPY_CASTING
 string_to_string_resolve_descriptors(PyObject *NPY_UNUSED(self),
                                      PyArray_DTypeMeta *NPY_UNUSED(dtypes[2]),
@@ -67,23 +69,15 @@ string_to_string(PyArrayMethod_Context *NPY_UNUSED(context),
     return 0;
 }
 
-static PyArray_DTypeMeta *s2s_dtypes[2] = {NULL, NULL};
-
 static PyType_Slot s2s_slots[] = {
         {NPY_METH_resolve_descriptors, &string_to_string_resolve_descriptors},
         {NPY_METH_strided_loop, &string_to_string},
         {NPY_METH_unaligned_strided_loop, &string_to_string},
         {0, NULL}};
 
-PyArrayMethod_Spec StringToStringCastSpec = {
-        .name = "cast_StringDType_to_StringDType",
-        .nin = 1,
-        .nout = 1,
-        .casting = NPY_NO_CASTING,
-        .flags = NPY_METH_SUPPORTS_UNALIGNED,
-        .dtypes = s2s_dtypes,
-        .slots = s2s_slots,
-};
+static char *s2s_name = "cast_StringDType_to_StringDType";
+
+// unicode to string
 
 static NPY_CASTING
 unicode_to_string_resolve_descriptors(PyObject *NPY_UNUSED(self),
@@ -261,6 +255,8 @@ static PyType_Slot u2s_slots[] = {
 
 static char *u2s_name = "cast_Unicode_to_StringDType";
 
+// string to unicode
+
 static NPY_CASTING
 string_to_unicode_resolve_descriptors(PyObject *NPY_UNUSED(self),
                                       PyArray_DTypeMeta *NPY_UNUSED(dtypes[2]),
@@ -376,44 +372,128 @@ static PyType_Slot s2u_slots[] = {
 
 static char *s2u_name = "cast_StringDType_to_Unicode";
 
+// string to bool
+
+static NPY_CASTING
+string_to_bool_resolve_descriptors(PyObject *NPY_UNUSED(self),
+                                   PyArray_DTypeMeta *NPY_UNUSED(dtypes[2]),
+                                   PyArray_Descr *given_descrs[2],
+                                   PyArray_Descr *loop_descrs[2],
+                                   npy_intp *NPY_UNUSED(view_offset))
+{
+    if (given_descrs[1] == NULL) {
+        loop_descrs[1] = PyArray_DescrNewFromType(NPY_BOOL);
+    }
+    else {
+        Py_INCREF(given_descrs[1]);
+        loop_descrs[1] = given_descrs[1];
+    }
+
+    Py_INCREF(given_descrs[0]);
+    loop_descrs[0] = given_descrs[0];
+
+    return NPY_UNSAFE_CASTING;
+}
+
+static int
+string_to_bool(PyArrayMethod_Context *context, char *const data[],
+               npy_intp const dimensions[], npy_intp const strides[],
+               NpyAuxData *NPY_UNUSED(auxdata))
+{
+    npy_intp N = dimensions[0];
+    char *in = data[0];
+    char *out = data[1];
+
+    npy_intp in_stride = strides[0];
+    npy_intp out_stride = strides[1];
+
+    ss *s = NULL;
+
+    while (N--) {
+        load_string(in, &s);
+        if (s->len == 0) {
+            *out = (npy_bool)0;
+        }
+        else {
+            *out = (npy_bool)1;
+        }
+
+        in += in_stride;
+        out += out_stride;
+    }
+
+    return 0;
+}
+
+static PyType_Slot s2b_slots[] = {
+        {NPY_METH_resolve_descriptors, &string_to_bool_resolve_descriptors},
+        {NPY_METH_strided_loop, &string_to_bool},
+        {0, NULL}};
+
+static char *s2b_name = "cast_StringDType_to_Bool";
+
+PyArrayMethod_Spec *
+get_cast_spec(const char *name, NPY_CASTING casting,
+              NPY_ARRAYMETHOD_FLAGS flags, PyArray_DTypeMeta **dtypes,
+              PyType_Slot *slots)
+{
+    PyArrayMethod_Spec *ret = malloc(sizeof(PyArrayMethod_Spec));
+
+    ret->name = name;
+    ret->nin = 1;
+    ret->nout = 1;
+    ret->casting = casting;
+    ret->flags = flags;
+    ret->dtypes = dtypes;
+    ret->slots = slots;
+
+    return ret;
+}
+
+PyArray_DTypeMeta **
+get_dtypes(PyArray_DTypeMeta *dt1, PyArray_DTypeMeta *dt2)
+{
+    PyArray_DTypeMeta **ret = malloc(2 * sizeof(PyArray_DTypeMeta *));
+
+    ret[0] = dt1;
+    ret[1] = dt2;
+
+    return ret;
+}
+
 PyArrayMethod_Spec **
 get_casts(void)
 {
-    PyArray_DTypeMeta **u2s_dtypes = malloc(2 * sizeof(PyArray_DTypeMeta *));
-    u2s_dtypes[0] = &PyArray_UnicodeDType;
-    u2s_dtypes[1] = NULL;
-
-    PyArrayMethod_Spec *UnicodeToStringCastSpec =
-            malloc(sizeof(PyArrayMethod_Spec));
-
-    UnicodeToStringCastSpec->name = u2s_name;
-    UnicodeToStringCastSpec->nin = 1;
-    UnicodeToStringCastSpec->nout = 1;
-    UnicodeToStringCastSpec->casting = NPY_SAFE_CASTING;
-    UnicodeToStringCastSpec->flags = NPY_METH_NO_FLOATINGPOINT_ERRORS;
-    UnicodeToStringCastSpec->dtypes = u2s_dtypes;
-    UnicodeToStringCastSpec->slots = u2s_slots;
-
-    PyArray_DTypeMeta **s2u_dtypes = malloc(2 * sizeof(PyArray_DTypeMeta *));
-    s2u_dtypes[0] = NULL;
-    s2u_dtypes[1] = &PyArray_UnicodeDType;
-
-    PyArrayMethod_Spec *StringToUnicodeCastSpec =
-            malloc(sizeof(PyArrayMethod_Spec));
-
-    StringToUnicodeCastSpec->name = s2u_name;
-    StringToUnicodeCastSpec->nin = 1;
-    StringToUnicodeCastSpec->nout = 1;
-    StringToUnicodeCastSpec->casting = NPY_SAFE_CASTING;
-    StringToUnicodeCastSpec->flags = NPY_METH_NO_FLOATINGPOINT_ERRORS;
-    StringToUnicodeCastSpec->dtypes = s2u_dtypes;
-    StringToUnicodeCastSpec->slots = s2u_slots;
-
-    PyArrayMethod_Spec **casts = malloc(4 * sizeof(PyArrayMethod_Spec *));
-    casts[0] = &StringToStringCastSpec;
+    PyArray_DTypeMeta **s2s_dtypes = get_dtypes(NULL, NULL);
+
+    PyArrayMethod_Spec *StringToStringCastSpec =
+            get_cast_spec(s2s_name, NPY_NO_CASTING,
+                          NPY_METH_SUPPORTS_UNALIGNED, s2s_dtypes, s2s_slots);
+
+    PyArray_DTypeMeta **u2s_dtypes = get_dtypes(&PyArray_UnicodeDType, NULL);
+
+    PyArrayMethod_Spec *UnicodeToStringCastSpec = get_cast_spec(
+            u2s_name, NPY_SAFE_CASTING, NPY_METH_NO_FLOATINGPOINT_ERRORS,
+            u2s_dtypes, u2s_slots);
+
+    PyArray_DTypeMeta **s2u_dtypes = get_dtypes(NULL, &PyArray_UnicodeDType);
+
+    PyArrayMethod_Spec *StringToUnicodeCastSpec = get_cast_spec(
+            s2u_name, NPY_SAFE_CASTING, NPY_METH_NO_FLOATINGPOINT_ERRORS,
+            s2u_dtypes, s2u_slots);
+
+    PyArray_DTypeMeta **s2b_dtypes = get_dtypes(NULL, &PyArray_BoolDType);
+
+    PyArrayMethod_Spec *StringToBoolCastSpec = get_cast_spec(
+            s2b_name, NPY_UNSAFE_CASTING, NPY_METH_NO_FLOATINGPOINT_ERRORS,
+            s2b_dtypes, s2b_slots);
+
+    PyArrayMethod_Spec **casts = malloc(5 * sizeof(PyArrayMethod_Spec *));
+    casts[0] = StringToStringCastSpec;
     casts[1] = UnicodeToStringCastSpec;
     casts[2] = StringToUnicodeCastSpec;
-    casts[3] = NULL;
+    casts[3] = StringToBoolCastSpec;
+    casts[4] = NULL;
 
     return casts;
 }

stringdtype/stringdtype/src/dtype.c

@@ -409,11 +409,10 @@ init_string_dtype(void)
 
     StringDType.singleton = singleton;
 
-    free(StringDType_DTypeSpec.casts[1]->dtypes);
-    free(StringDType_DTypeSpec.casts[1]);
-    free(StringDType_DTypeSpec.casts[2]->dtypes);
-    free(StringDType_DTypeSpec.casts[2]);
-    free(StringDType_DTypeSpec.casts);
+    for (int i = 0; casts[i] != NULL; i++) {
+        free(casts[i]->dtypes);
+        free(casts[i]);
+    }
 
     return 0;
 }

stringdtype/tests/test_stringdtype.py

@@ -242,3 +242,20 @@ def test_arrfuncs_empty(arrfunc, expected):
     arr = np.empty(10, dtype=StringDType())
     result = arrfunc(arr)
     np.testing.assert_array_equal(result, expected, strict=True)
+
+    
+@pytest.mark.parametrize(
+    ("string_list", "cast_answer", "any_answer", "all_answer"),
+    [
+        [["hello", "world"], [True, True], True, True],
+        [["", ""], [False, False], False, False],
+        [["hello", ""], [True, False], True, False],
+        [["", "world"], [False, True], True, False],
+    ],
+)
+def test_bool_cast(string_list, cast_answer, any_answer, all_answer):
+    sarr = np.array(string_list, dtype=StringDType())
+    np.testing.assert_array_equal(sarr.astype("bool"), cast_answer)
+
+    assert np.any(sarr) == any_answer
+    assert np.all(sarr) == all_answer