WIP: add casts between stringdtype and pandasstringdtype

Author: ngoldbaum

stringdtype/stringdtype/src/casts.c

@@ -75,6 +75,8 @@ static PyType_Slot s2s_slots[] = {
         {0, NULL}};
 
 static char *s2s_name = "cast_StringDType_to_StringDType";
+static char *s2p_name = "cast_StringDType_to_PandasStringDType";
+static char *p2s_name = "cast_PandasStringDType_to_StringDType";
 
 // unicode to string
 
@@ -476,38 +478,83 @@ get_dtypes(PyArray_DTypeMeta *dt1, PyArray_DTypeMeta *dt2)
 }
 
 PyArrayMethod_Spec **
-get_casts(void)
+get_casts(PyArray_DTypeMeta *this, PyArray_DTypeMeta *other,
+          int pandas_available)
 {
-    PyArray_DTypeMeta **s2s_dtypes = get_dtypes(NULL, NULL);
+    PyArray_DTypeMeta **s2s_dtypes = get_dtypes(this, this);
 
     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 *ThisToOtherCastSpec = NULL;
+    PyArrayMethod_Spec *OtherToThisCastSpec = NULL;
+
+    if (pandas_available) {
+        char *t2o_name = NULL;
+        char *o2t_name = NULL;
+
+        if (this == (PyArray_DTypeMeta *)&StringDType) {
+            t2o_name = s2p_name;
+            o2t_name = p2s_name;
+        }
+        else {
+            t2o_name = p2s_name;
+            o2t_name = s2p_name;
+        }
+
+        PyArray_DTypeMeta **t2o_dtypes = get_dtypes(this, other);
+
+        ThisToOtherCastSpec = get_cast_spec(t2o_name, NPY_NO_CASTING,
+                                            NPY_METH_SUPPORTS_UNALIGNED,
+                                            t2o_dtypes, s2s_slots);
+
+        PyArray_DTypeMeta **o2t_dtypes = get_dtypes(other, this);
+
+        OtherToThisCastSpec = get_cast_spec(o2t_name, NPY_NO_CASTING,
+                                            NPY_METH_SUPPORTS_UNALIGNED,
+                                            o2t_dtypes, s2s_slots);
+    }
+
+    PyArray_DTypeMeta **u2s_dtypes = get_dtypes(&PyArray_UnicodeDType, this);
 
     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);
+    PyArray_DTypeMeta **s2u_dtypes = get_dtypes(this, &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);
+    PyArray_DTypeMeta **s2b_dtypes = get_dtypes(this, &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 *));
+    PyArrayMethod_Spec **casts = NULL;
+
+    if (pandas_available) {
+        casts = malloc(7 * sizeof(PyArrayMethod_Spec *));
+    }
+    else {
+        casts = malloc(5 * sizeof(PyArrayMethod_Spec *));
+    }
+
     casts[0] = StringToStringCastSpec;
     casts[1] = UnicodeToStringCastSpec;
     casts[2] = StringToUnicodeCastSpec;
     casts[3] = StringToBoolCastSpec;
-    casts[4] = NULL;
+    if (pandas_available) {
+        casts[4] = ThisToOtherCastSpec;
+        casts[5] = OtherToThisCastSpec;
+        casts[6] = NULL;
+    }
+    else {
+        casts[4] = NULL;
+    }
 
     return casts;
 }

stringdtype/stringdtype/src/casts.h

@@ -11,7 +11,8 @@
 #include "numpy/ndarraytypes.h"
 
 PyArrayMethod_Spec **
-get_casts(void);
+get_casts(PyArray_DTypeMeta *this_dtype, PyArray_DTypeMeta *other_dtype,
+          int pandas_available);
 
 size_t
 utf8_char_to_ucs4_code(unsigned char *, Py_UCS4 *);

stringdtype/stringdtype/src/dtype.c

@@ -523,12 +523,24 @@ StringDType_type PandasStringDType = {
 int
 init_string_dtype(void)
 {
-    PyArrayMethod_Spec **casts = get_casts();
+    PyObject *pandas_mod = PyImport_ImportModule("pandas");
+
+    if (pandas_mod == NULL) {
+        // clear ImportError
+        PyErr_Clear();
+    }
+    else {
+        PANDAS_AVAILABLE = 1;
+    }
+
+    PyArrayMethod_Spec **StringDType_casts = get_casts(
+            (PyArray_DTypeMeta *)&StringDType,
+            (PyArray_DTypeMeta *)&PandasStringDType, PANDAS_AVAILABLE);
 
     PyArrayDTypeMeta_Spec StringDType_DTypeSpec = {
             .typeobj = StringScalar_Type,
             .slots = StringDType_Slots,
-            .casts = casts,
+            .casts = StringDType_casts,
     };
 
     /* Loaded dynamically, so may need to be set here: */
@@ -545,6 +557,10 @@ init_string_dtype(void)
         return -1;
     }
 
+    // Partially initialize PandasStringDType so cast setup succeeds
+    ((PyObject *)&PandasStringDType)->ob_type = &PyArrayDTypeMeta_Type;
+    ((PyTypeObject *)&PandasStringDType)->tp_base = &PyArrayDescr_Type;
+
     if (PyArrayInitDTypeMeta_FromSpec((PyArray_DTypeMeta *)&StringDType,
                                       &StringDType_DTypeSpec) < 0) {
         return -1;
@@ -559,27 +575,32 @@ init_string_dtype(void)
 
     StringDType.base.singleton = singleton;
 
+    for (int i = 0; StringDType_casts[i] != NULL; i++) {
+        free(StringDType_casts[i]->dtypes);
+        free(StringDType_casts[i]);
+    }
+
     /* and once again for PandasStringDType */
 
-    PyObject *mod = PyImport_ImportModule("pandas");
+    if (PANDAS_AVAILABLE) {
+        PyArrayMethod_Spec **PandasStringDType_casts =
+                get_casts((PyArray_DTypeMeta *)&PandasStringDType,
+                          (PyArray_DTypeMeta *)&StringDType, PANDAS_AVAILABLE);
 
-    if (mod != NULL) {
         PyArrayDTypeMeta_Spec PandasStringDType_DTypeSpec = {
                 .typeobj = PandasStringScalar_Type,
                 .slots = StringDType_Slots,
-                .casts = casts,
+                .casts = PandasStringDType_casts,
         };
 
-        PyObject *pandas_na_obj = PyObject_GetAttrString(mod, "NA");
+        PyObject *pandas_na_obj = PyObject_GetAttrString(pandas_mod, "NA");
 
-        Py_DECREF(mod);
+        Py_DECREF(pandas_mod);
 
         if (pandas_na_obj == NULL) {
             return -1;
         }
 
-        ((PyObject *)&PandasStringDType)->ob_type = &PyArrayDTypeMeta_Type;
-        ((PyTypeObject *)&PandasStringDType)->tp_base = &PyArrayDescr_Type;
         ((PyTypeObject *)&PandasStringDType)->tp_dict = PyDict_New();
         // C attribute for fast access
         Py_INCREF(pandas_na_obj);
@@ -605,15 +626,11 @@ init_string_dtype(void)
         }
 
         PandasStringDType.base.singleton = singleton;
-        PANDAS_AVAILABLE = 1;
-    }
-    else {
-        PyErr_Clear();
-    }
 
-    for (int i = 0; casts[i] != NULL; i++) {
-        free(casts[i]->dtypes);
-        free(casts[i]);
+        for (int i = 0; PandasStringDType_casts[i] != NULL; i++) {
+            free(PandasStringDType_casts[i]->dtypes);
+            free(PandasStringDType_casts[i]);
+        }
     }
 
     return 0;