initialize stringdtype instances using a generic type

Author: ngoldbaum

stringdtype/stringdtype/src/casts.c

@@ -80,17 +80,18 @@ static char *s2s_name = "cast_StringDType_to_StringDType";
 
 static NPY_CASTING
 unicode_to_string_resolve_descriptors(PyObject *NPY_UNUSED(self),
-                                      PyArray_DTypeMeta *NPY_UNUSED(dtypes[2]),
+                                      PyArray_DTypeMeta *dtypes[2],
                                       PyArray_Descr *given_descrs[2],
                                       PyArray_Descr *loop_descrs[2],
                                       npy_intp *NPY_UNUSED(view_offset))
 {
     if (given_descrs[1] == NULL) {
-        StringDTypeObject *new = new_stringdtype_instance();
+        PyArray_Descr *new = (PyArray_Descr *)new_stringdtype_instance(
+                (PyTypeObject *)dtypes[1]);
         if (new == NULL) {
             return (NPY_CASTING)-1;
         }
-        loop_descrs[1] = (PyArray_Descr *)new;
+        loop_descrs[1] = new;
     }
     else {
         Py_INCREF(given_descrs[1]);

stringdtype/stringdtype/src/dtype.c

@@ -10,19 +10,20 @@ PyObject *NA_OBJ = NULL;
 /*
  * Internal helper to create new instances
  */
-StringDTypeObject *
-new_stringdtype_instance(void)
+PyObject *
+new_stringdtype_instance(PyTypeObject *cls)
 {
-    StringDTypeObject *new = (StringDTypeObject *)PyArrayDescr_Type.tp_new(
-            (PyTypeObject *)&StringDType, NULL, NULL);
+    PyObject *new = PyArrayDescr_Type.tp_new((PyTypeObject *)cls, NULL, NULL);
     if (new == NULL) {
         return NULL;
     }
-    new->base.elsize = sizeof(ss);
-    new->base.alignment = _Alignof(ss);
-    new->base.flags |= NPY_NEEDS_INIT;
-    new->base.flags |= NPY_LIST_PICKLE;
-    new->base.flags |= NPY_ITEM_REFCOUNT;
+
+    PyArray_Descr *base = &((StringDTypeObject *)new)->base;
+    base->elsize = sizeof(ss);
+    base->alignment = _Alignof(ss);
+    base->flags |= NPY_NEEDS_INIT;
+    base->flags |= NPY_LIST_PICKLE;
+    base->flags |= NPY_ITEM_REFCOUNT;
 
     return new;
 }
@@ -63,16 +64,15 @@ common_dtype(PyArray_DTypeMeta *cls, PyArray_DTypeMeta *other)
 // For a given python object, this function returns a borrowed reference
 // to the dtype property of the array
 static PyArray_Descr *
-string_discover_descriptor_from_pyobject(PyArray_DTypeMeta *NPY_UNUSED(cls),
-                                         PyObject *obj)
+string_discover_descriptor_from_pyobject(PyTypeObject *cls, PyObject *obj)
 {
     if (Py_TYPE(obj) != StringScalar_Type) {
         PyErr_SetString(PyExc_TypeError,
                         "Can only store StringScalar in a StringDType array.");
         return NULL;
     }
 
-    PyArray_Descr *ret = (PyArray_Descr *)new_stringdtype_instance();
+    PyArray_Descr *ret = (PyArray_Descr *)new_stringdtype_instance(cls);
     if (ret == NULL) {
         return NULL;
     }
@@ -354,7 +354,7 @@ static PyType_Slot StringDType_Slots[] = {
         {0, NULL}};
 
 static PyObject *
-stringdtype_new(PyTypeObject *NPY_UNUSED(cls), PyObject *args, PyObject *kwds)
+stringdtype_new(PyTypeObject *cls, PyObject *args, PyObject *kwds)
 {
     static char *kwargs_strs[] = {"size", NULL};
 
@@ -365,7 +365,7 @@ stringdtype_new(PyTypeObject *NPY_UNUSED(cls), PyObject *args, PyObject *kwds)
         return NULL;
     }
 
-    PyObject *ret = (PyObject *)new_stringdtype_instance();
+    PyObject *ret = new_stringdtype_instance(cls);
 
     return ret;
 }
@@ -494,7 +494,6 @@ init_string_dtype(void)
     PyArrayMethod_Spec **casts = get_casts();
 
     PyArrayDTypeMeta_Spec StringDType_DTypeSpec = {
-            .flags = NPY_DT_PARAMETRIC,
             .typeobj = StringScalar_Type,
             .slots = StringDType_Slots,
             .casts = casts,

stringdtype/stringdtype/src/dtype.h

@@ -16,15 +16,14 @@
 
 typedef struct {
     PyArray_Descr base;
-    PyObject *na_object;
 } StringDTypeObject;
 
 extern PyArray_DTypeMeta StringDType;
 extern PyTypeObject *StringScalar_Type;
 extern PyObject *NA_OBJ;
 
-StringDTypeObject *
-new_stringdtype_instance(void);
+PyObject *
+new_stringdtype_instance(PyTypeObject *cls);
 
 int
 init_string_dtype(void);

stringdtype/stringdtype/src/umath.c

@@ -15,21 +15,19 @@
 
 static NPY_CASTING
 binary_resolve_descriptors(struct PyArrayMethodObject_tag *NPY_UNUSED(method),
-                           PyArray_DTypeMeta *NPY_UNUSED(dtypes[]),
+                           PyArray_DTypeMeta *dtypes[],
                            PyArray_Descr *given_descrs[],
                            PyArray_Descr *loop_descrs[],
                            npy_intp *NPY_UNUSED(view_offset))
 {
-    PyObject *na_obj1 = ((StringDTypeObject *)given_descrs[0])->na_object;
-    PyObject *na_obj2 = ((StringDTypeObject *)given_descrs[1])->na_object;
-
-    int eq_res = PyObject_RichCompareBool(na_obj1, na_obj2, Py_EQ);
-
-    if (eq_res < 0) {
-        return (NPY_CASTING)-1;
-    }
-
-    if (eq_res != 1) {
+    // technically incorrect to cast to StringDType if we have a
+    // PandasStringDType but they have the same layout so this should be fine.
+    PyObject *na_obj1 = PyDict_GetItemString(
+            ((PyTypeObject *)dtypes[0])->tp_dict, "na_object");
+    PyObject *na_obj2 = PyDict_GetItemString(
+            ((PyTypeObject *)dtypes[1])->tp_dict, "na_object");
+
+    if (na_obj1 != na_obj2) {
         PyErr_SetString(PyExc_TypeError,
                         "Can only do binary operations with identical "
                         "StringDType instances.");