fixes for StringDType

Author: ngoldbaum

asciidtype/asciidtype/scalar.py

@@ -6,3 +6,11 @@ def __new__(cls, value, dtype):
         instance = super().__new__(cls, value)
         instance.dtype = dtype
         return instance
+
+    def partition(self, sep):
+        ret = super().partition(sep)
+        return (str(ret[0]), str(ret[1]), str(ret[2]))
+
+    def rpartition(self, sep):
+        ret = super().rpartition(sep)
+        return (str(ret[0]), str(ret[1]), str(ret[2]))

asciidtype/asciidtype/src/dtype.c

@@ -88,15 +88,31 @@ static PyArray_Descr *
 ascii_discover_descriptor_from_pyobject(PyArray_DTypeMeta *NPY_UNUSED(cls),
                                         PyObject *obj)
 {
-    if (Py_TYPE(obj) != ASCIIScalar_Type) {
+    PyTypeObject *obj_type = Py_TYPE(obj);
+    PyArray_Descr *ret = NULL;
+    if (obj_type != ASCIIScalar_Type) {
+        if (PyUnicode_Check(obj)) {
+            if (!PyUnicode_IS_ASCII(obj)) {
+                PyErr_SetString(
+                        PyExc_TypeError,
+                        "Can only store strings or bytes convertible to ASCII "
+                        "in a ASCIIDType array.");
+                return NULL;
+            }
+            ret = (PyArray_Descr *)new_asciidtype_instance(
+                    (long)PyUnicode_GetLength(obj));
+        }
+        // could do bytes too if we want
         PyErr_SetString(PyExc_TypeError,
-                        "Can only store ASCIIScalar in a ASCIIDType array.");
+                        "Can only store strings or bytes convertible to ASCII "
+                        "in a ASCIIDType array.");
         return NULL;
     }
-
-    PyArray_Descr *ret = (PyArray_Descr *)PyObject_GetAttrString(obj, "dtype");
-    if (ret == NULL) {
-        return NULL;
+    else {
+        ret = (PyArray_Descr *)PyObject_GetAttrString(obj, "dtype");
+        if (ret == NULL) {
+            return NULL;
+        }
     }
     return ret;
 }

stringdtype/README.md

@@ -1,6 +1,6 @@
-# A dtype that stores ASCII data
+# A dtype that stores pointers to strings
 
-This is a simple proof-of-concept dtype using the (as of late 2022) experimental
+This is a simple proof-of-concept dtype using the (as of early 2023) experimental
 [new dtype
 implementation](https://numpy.org/neps/nep-0041-improved-dtype-support.html) in
 NumPy.

stringdtype/stringdtype/scalar.py

@@ -1,12 +1,16 @@
 """A scalar type needed by the dtype machinery."""
 
 
-class StringScalar:
-    def __init__(self, value):
-        self.value = value
+class StringScalar(str):
+    def __new__(cls, value, dtype):
+        instance = super().__new__(cls, value)
+        instance.dtype = dtype
+        return instance
 
-    def __str__(self):
-        return str(self.value)
+    def partition(self, sep):
+        ret = super().partition(sep)
+        return (str(ret[0]), str(ret[1]), str(ret[2]))
 
-    def __repr__(self):
-        return repr(self.value)
+    def rpartition(self, sep):
+        ret = super().rpartition(sep)
+        return (str(ret[0]), str(ret[1]), str(ret[2]))

stringdtype/stringdtype/src/dtype.c

@@ -21,20 +21,16 @@ new_stringdtype_instance(void)
     return new;
 }
 
-//
-// This is used to determine the correct dtype to return when operations mix
-// dtypes (I think?). For now just return the first one.
-//
+/*
+ * This is used to determine the correct dtype to return when dealing
+ * with a mix of different dtypes (for example when creating an array
+ * from a list of scalars). Since StringDType doesn't have any parameters,
+ * we can safely always return the first one.
+ */
 static StringDTypeObject *
 common_instance(StringDTypeObject *dtype1, StringDTypeObject *dtype2)
 {
-    if (!PyObject_RichCompareBool((PyObject *)dtype1, (PyObject *)dtype2,
-                                  Py_EQ)) {
-        PyErr_SetString(
-                PyExc_RuntimeError,
-                "common_instance called on unequal StringDType instances");
-        return NULL;
-    }
+    Py_INCREF(dtype1);
     return dtype1;
 }
 
@@ -66,34 +62,65 @@ string_discover_descriptor_from_pyobject(PyArray_DTypeMeta *NPY_UNUSED(cls),
     return ret;
 }
 
+static PyObject *
+get_value(PyObject *scalar)
+{
+    PyObject *ret_bytes = NULL;
+    PyTypeObject *scalar_type = Py_TYPE(scalar);
+    // FIXME: handle bytes too
+    if ((scalar_type == &PyUnicode_Type) ||
+        (scalar_type == StringScalar_Type)) {
+        // attempt to decode as UTF8
+        ret_bytes = PyUnicode_AsUTF8String(scalar);
+        if (ret_bytes == NULL) {
+            PyErr_SetString(
+                    PyExc_TypeError,
+                    "Can only store UTF8 text in a StringDType array.");
+            return NULL;
+        }
+    }
+    else {
+        PyErr_SetString(PyExc_TypeError,
+                        "Can only store String text in a StringDType array.");
+        return NULL;
+    }
+    return ret_bytes;
+}
+
 // Take a python object `obj` and insert it into the array of dtype `descr` at
 // the position given by dataptr.
 static int
 stringdtype_setitem(StringDTypeObject *descr, PyObject *obj, char **dataptr)
 {
-    char *val = PyBytes_AsString(obj);
-    if (val == NULL) {
+    PyObject *val_obj = get_value(obj);
+    char *val = NULL;
+    Py_ssize_t length = 0;
+    if (PyBytes_AsStringAndSize(val_obj, &val, &length) == -1) {
         return -1;
     }
 
-    *dataptr = malloc(sizeof(char) * strlen(val));
-    strcpy(*dataptr, val);
+    *dataptr = malloc(sizeof(char) * length + 1);
+    strncpy(*dataptr, val, length + 1);
+    Py_DECREF(val_obj);
     return 0;
 }
 
 static PyObject *
 stringdtype_getitem(StringDTypeObject *descr, char **dataptr)
 {
-    PyObject *val_obj = PyBytes_FromString(*dataptr);
+    PyObject *val_obj = PyUnicode_FromString(*dataptr);
+
     if (val_obj == NULL) {
         return NULL;
     }
 
     PyObject *res = PyObject_CallFunctionObjArgs((PyObject *)StringScalar_Type,
-                                                 val_obj, NULL);
+                                                 val_obj, descr, NULL);
+
     if (res == NULL) {
         return NULL;
     }
+
     Py_DECREF(val_obj);
 
     return res;
@@ -119,6 +146,15 @@ static PyType_Slot StringDType_Slots[] = {
 static PyObject *
 stringdtype_new(PyTypeObject *NPY_UNUSED(cls), PyObject *args, PyObject *kwds)
 {
+    static char *kwargs_strs[] = {"size", NULL};
+
+    long size = 0;
+
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|l:StringDType", kwargs_strs,
+                                     &size)) {
+        return NULL;
+    }
+
     return (PyObject *)new_stringdtype_instance();
 }