add unicode_ to StringDType cast

ngoldbaum · ngoldbaum · commit bc3f8bed8696 · 2023-01-16T11:55:24.000-07:00
diff --git a/stringdtype/stringdtype/src/casts.c b/stringdtype/stringdtype/src/casts.c
@@ -1,18 +1,19 @@
 #include "casts.h"
 
+#include "dtype.h"
+
 static NPY_CASTING
-string_to_string_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))
+string_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))
 {
     Py_INCREF(given_descrs[0]);
     loop_descrs[0] = given_descrs[0];
 
     if (given_descrs[1] == NULL) {
-        Py_INCREF(given_descrs[0]);
-        loop_descrs[1] = given_descrs[0];
+        loop_descrs[1] = (PyArray_Descr *)new_stringdtype_instance();
     }
     else {
         Py_INCREF(given_descrs[1]);
@@ -49,7 +50,7 @@ string_to_string(PyArrayMethod_Context *context, char *const data[],
 static PyArray_DTypeMeta *s2s_dtypes[2] = {NULL, NULL};
 
 static PyType_Slot s2s_slots[] = {
-        {NPY_METH_resolve_descriptors, &string_to_string_resolve_descriptors},
+        {NPY_METH_resolve_descriptors, &string_resolve_descriptors},
         {NPY_METH_strided_loop, &string_to_string},
         {NPY_METH_unaligned_strided_loop, &string_to_string},
         {0, NULL}};
@@ -64,12 +65,137 @@ PyArrayMethod_Spec StringToStringCastSpec = {
         .slots = s2s_slots,
 };
 
+// converts UCS4 code point to 4-byte char* assumes in is a zero-filled 4 byte
+// array returns -1 if the code point is not a valid unicode code point, the
+// number of bytes in the in the UTF-8 character on success
+static int
+ucs4_to_utf8_char(const Py_UCS4 code, char *in)
+{
+    if (code <= 0x7F) {
+        // ASCII
+        in[0] = (char)code;
+        return 1;
+    }
+    else if (code <= 0x07FF) {
+        in[0] = (0xc0 | (code >> 6));
+        in[1] = (0x80 | (code & 0x3f));
+        return 2;
+    }
+    else if (code <= 0xFFFF) {
+        in[0] = (0xe0 | (code >> 12));
+        in[1] = (0x80 | ((code >> 6) & 0x3f));
+        in[2] = (0x80 | (code & 0x3f));
+        return 3;
+    }
+    else if (code <= 0x10FFFF) {
+        in[0] = (0xf0 | (code >> 18));
+        in[1] = (0x80 | ((code >> 12) & 0x3f));
+        in[2] = (0x80 | ((code >> 6) & 0x3f));
+        in[3] = (0x80 | (code & 0x3f));
+        return 4;
+    }
+    return -1;
+}
+
+static int
+unicode_to_string(PyArrayMethod_Context *context, char *const data[],
+                  npy_intp const dimensions[], npy_intp const strides[],
+                  NpyAuxData *NPY_UNUSED(auxdata))
+{
+    PyArray_Descr **descrs = context->descriptors;
+    long in_size = (descrs[0]->elsize) / 4;
+
+    npy_intp N = dimensions[0];
+    char *in = data[0];
+    char **out = (char **)data[1];
+    npy_intp in_stride = strides[0];
+    // strides are in bytes but pointer offsets are in pointer widths, so
+    // divide by the element size (one pointer width) to get the pointer offset
+    npy_intp out_stride = strides[1] / context->descriptors[1]->elsize;
+
+    while (N--) {
+        // pessimistically allocate 4 bytes per allowed character
+        char *out_buf = calloc(in_size * 4 + 1, sizeof(char));
+        size_t out_num_bytes = 0;
+        for (int i = 0; i < in_size; i++) {
+            // get code point
+            Py_UCS4 code = ((Py_UCS4 *)in)[i];
+
+            if (code == 0) {
+                break;
+            }
+
+            // convert codepoint to UTF8 bytes
+            char utf8_c[4] = {0};
+            size_t num_bytes = ucs4_to_utf8_char(code, utf8_c);
+            out_num_bytes += num_bytes;
+
+            if (num_bytes == -1) {
+                // acquire GIL, set error, return
+                PyGILState_STATE gstate;
+                gstate = PyGILState_Ensure();
+                PyErr_SetString(PyExc_TypeError,
+                                "Invalid unicode code point found");
+                PyGILState_Release(gstate);
+                return -1;
+            }
+
+            // copy utf8_c into out_buf
+            strncpy(out_buf, utf8_c, num_bytes);
+
+            // increment out_buf by the size of the character
+            out_buf += num_bytes;
+        }
+
+        // reset out_buf to the beginning of the string
+        out_buf -= out_num_bytes;
+
+        // pad string with null character
+        out_buf[out_num_bytes] = '\0';
+
+        // resize out_buf now that we know the real size
+        out_buf = realloc(out_buf, out_num_bytes + 1);
+
+        // set out to the address of the beginning of the string
+        out[0] = out_buf;
+
+        // increment out and in by strides
+        in += in_stride;
+        out += out_stride;
+    }
+
+    return 0;
+}
+
+static PyType_Slot u2s_slots[] = {
+        {NPY_METH_resolve_descriptors, &string_resolve_descriptors},
+        {NPY_METH_strided_loop, &unicode_to_string},
+        {0, NULL}};
+
+static char *u2s_name = "cast_Unicode_to_StringDType";
+
 PyArrayMethod_Spec **
 get_casts(void)
 {
-    PyArrayMethod_Spec **casts = malloc(2 * sizeof(PyArrayMethod_Spec *));
+    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;
+
+    PyArrayMethod_Spec **casts = malloc(3 * sizeof(PyArrayMethod_Spec *));
     casts[0] = &StringToStringCastSpec;
-    casts[1] = NULL;
+    casts[1] = UnicodeToStringCastSpec;
+    casts[2] = NULL;
 
     return casts;
 }
diff --git a/stringdtype/stringdtype/src/dtype.c b/stringdtype/stringdtype/src/dtype.c
@@ -226,5 +226,9 @@ init_string_dtype(void)
 
     StringDType.singleton = singleton;
 
+    free(StringDType_DTypeSpec.casts[1]->dtypes);
+    free(StringDType_DTypeSpec.casts[1]);
+    free(StringDType_DTypeSpec.casts);
+
     return 0;
 }
diff --git a/stringdtype/tests/test_stringdtype.py b/stringdtype/tests/test_stringdtype.py
@@ -6,94 +6,99 @@
 
 @pytest.fixture
 def string_list():
-    return ['abc', 'def', 'ghi']
+    return ["abc", "def", "ghi"]
 
 
 def test_scalar_creation():
-    assert str(StringScalar('abc', StringDType())) == 'abc'
+    assert str(StringScalar("abc", StringDType())) == "abc"
 
 
 def test_dtype_creation():
-    assert str(StringDType()) == 'StringDType'
+    assert str(StringDType()) == "StringDType"
 
 
 @pytest.mark.parametrize(
-    'data', [
-        ['abc', 'def', 'ghi'],
+    "data",
+    [
+        ["abc", "def", "ghi"],
         ["🤣", "📵", "😰"],
         ["🚜", "🙃", "😾"],
         ["😹", "🚠", "🚌"],
-    ]
+    ],
 )
 def test_array_creation_utf8(data):
     arr = np.array(data, dtype=StringDType())
-    assert repr(arr) == f'array({str(data)}, dtype=StringDType)'
+    assert repr(arr) == f"array({str(data)}, dtype=StringDType)"
 
 
 def test_array_creation_scalars(string_list):
     dtype = StringDType()
     arr = np.array(
         [
-            StringScalar('abc', dtype=dtype),
-            StringScalar('def', dtype=dtype),
-            StringScalar('ghi', dtype=dtype),
+            StringScalar("abc", dtype=dtype),
+            StringScalar("def", dtype=dtype),
+            StringScalar("ghi", dtype=dtype),
         ]
     )
     assert repr(arr) == repr(np.array(string_list, dtype=StringDType()))
 
 
 @pytest.mark.parametrize(
-   'data', [
-       [1, 2, 3],
-       [None, None, None],
-       [b'abc', b'def', b'ghi'],
-       [object, object, object],
-   ]
+    "data",
+    [
+        [1, 2, 3],
+        [None, None, None],
+        [b"abc", b"def", b"ghi"],
+        [object, object, object],
+    ],
 )
 def test_bad_scalars(data):
     with pytest.raises(TypeError):
         np.array(data, dtype=StringDType())
 
 
-@pytest.mark.xfail(reason='Not yet implemented')
+@pytest.mark.parametrize(
+    ("string_list"),
+    [
+        ["this", "is", "an", "array"],
+        ["€", "", "😊"],
+    ],
+)
 def test_cast_to_stringdtype(string_list):
-    arr = np.array(string_list, dtype='<U3').astype(StringDType())
+    arr = np.array(string_list, dtype=np.unicode_).astype(StringDType())
     expected = np.array(string_list, dtype=StringDType())
     np.testing.assert_array_equal(arr, expected)
 
 
-@pytest.mark.xfail(reason='Not yet implemented')
+@pytest.mark.xfail(reason="Not yet implemented")
 def test_cast_to_unicode_safe(string_list):
     arr = np.array(string_list, dtype=StringDType())
 
     np.testing.assert_array_equal(
-        arr.astype('<U3', casting='safe'),
-        np.array(string_list, dtype='<U3')
+        arr.astype("<U3", casting="safe"), np.array(string_list, dtype="<U3")
     )
 
     # Safe casting should preserve data
     with pytest.raises(TypeError):
-        arr.astype('<U2', casting='safe')
+        arr.astype("<U2", casting="safe")
 
 
-@pytest.mark.xfail(reason='Not yet implemented')
+@pytest.mark.xfail(reason="Not yet implemented")
 def test_cast_to_unicode_unsafe(string_list):
     arr = np.array(string_list, dtype=StringDType())
 
     np.testing.assert_array_equal(
-        arr.astype('<U3', casting='unsafe'),
-        np.array(string_list, dtype='<U3')
+        arr.astype("<U3", casting="unsafe"), np.array(string_list, dtype="<U3")
     )
 
     # Unsafe casting: each element is truncated
     np.testing.assert_array_equal(
-        arr.astype('<U2', casting='unsafe'),
-        np.array(string_list, dtype='<U2')
+        arr.astype("<U2", casting="unsafe"), np.array(string_list, dtype="<U2")
     )
 
 
 def test_insert_scalar(string_list):
     dtype = StringDType()
     arr = np.array(string_list, dtype=dtype)
-    arr[1] = StringScalar('what', dtype=dtype)
-    assert repr(arr) == repr(np.array(['abc', 'what', 'ghi'], dtype=dtype))
+    arr[1] = StringScalar("what", dtype=dtype)
+    assert repr(arr) == repr(np.array(["abc", "what", "ghi"], dtype=dtype))

Original file line number	Diff line number	Diff line change
`@@ -226,5 +226,9 @@ init_string_dtype(void)`
`226`	`226`
`227`	`227`	`StringDType.singleton = singleton;`
`228`	`228`
	`229`	`+ free(StringDType_DTypeSpec.casts[1]->dtypes);`
	`230`	`+ free(StringDType_DTypeSpec.casts[1]);`
	`231`	`+ free(StringDType_DTypeSpec.casts);`
	`232`	`+`
`229`	`233`	`return 0;`
`230`	`234`	`}`