add round-trip casts between StringDType and PandasStringDType

Author: ngoldbaum

stringdtype/stringdtype/src/casts.c

@@ -75,6 +75,7 @@ static PyType_Slot s2s_slots[] = {
         {0, NULL}};
 
 static char *s2s_name = "cast_StringDType_to_StringDType";
+static char *p2p_name = "cast_PandasStringDType_to_PandasStringDType";
 static char *s2p_name = "cast_StringDType_to_PandasStringDType";
 static char *p2s_name = "cast_PandasStringDType_to_StringDType";
 
@@ -478,40 +479,42 @@ get_dtypes(PyArray_DTypeMeta *dt1, PyArray_DTypeMeta *dt2)
 }
 
 PyArrayMethod_Spec **
-get_casts(PyArray_DTypeMeta *this, PyArray_DTypeMeta *other,
-          int pandas_available)
+get_casts(PyArray_DTypeMeta *this, PyArray_DTypeMeta *other)
 {
-    PyArray_DTypeMeta **s2s_dtypes = get_dtypes(this, this);
+    char *t2t_name = NULL;
 
-    PyArrayMethod_Spec *StringToStringCastSpec =
-            get_cast_spec(s2s_name, NPY_NO_CASTING,
-                          NPY_METH_SUPPORTS_UNALIGNED, s2s_dtypes, s2s_slots);
+    if (this == (PyArray_DTypeMeta *)&StringDType) {
+        t2t_name = s2s_name;
+    }
+    else {
+        t2t_name = p2p_name;
+    }
+
+    PyArray_DTypeMeta **t2t_dtypes = get_dtypes(this, this);
+
+    PyArrayMethod_Spec *ThisToThisCastSpec =
+            get_cast_spec(t2t_name, NPY_NO_CASTING,
+                          NPY_METH_SUPPORTS_UNALIGNED, t2t_dtypes, s2s_slots);
 
     PyArrayMethod_Spec *ThisToOtherCastSpec = NULL;
     PyArrayMethod_Spec *OtherToThisCastSpec = NULL;
 
-    if (pandas_available) {
-        char *t2o_name = NULL;
-        char *o2t_name = NULL;
+    int is_pandas = (this == (PyArray_DTypeMeta *)&PandasStringDType);
 
-        if (this == (PyArray_DTypeMeta *)&StringDType) {
-            t2o_name = s2p_name;
-            o2t_name = p2s_name;
-        }
-        else {
-            t2o_name = p2s_name;
-            o2t_name = s2p_name;
-        }
+    int num_casts = 5;
+
+    if (is_pandas) {
+        num_casts = 7;
 
         PyArray_DTypeMeta **t2o_dtypes = get_dtypes(this, other);
 
-        ThisToOtherCastSpec = get_cast_spec(t2o_name, NPY_NO_CASTING,
+        ThisToOtherCastSpec = get_cast_spec(p2s_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,
+        OtherToThisCastSpec = get_cast_spec(s2p_name, NPY_NO_CASTING,
                                             NPY_METH_SUPPORTS_UNALIGNED,
                                             o2t_dtypes, s2s_slots);
     }
@@ -536,18 +539,13 @@ get_casts(PyArray_DTypeMeta *this, PyArray_DTypeMeta *other,
 
     PyArrayMethod_Spec **casts = NULL;
 
-    if (pandas_available) {
-        casts = malloc(7 * sizeof(PyArrayMethod_Spec *));
-    }
-    else {
-        casts = malloc(5 * sizeof(PyArrayMethod_Spec *));
-    }
+    casts = malloc(num_casts * sizeof(PyArrayMethod_Spec *));
 
-    casts[0] = StringToStringCastSpec;
+    casts[0] = ThisToThisCastSpec;
     casts[1] = UnicodeToStringCastSpec;
     casts[2] = StringToUnicodeCastSpec;
     casts[3] = StringToBoolCastSpec;
-    if (pandas_available) {
+    if (is_pandas) {
         casts[4] = ThisToOtherCastSpec;
         casts[5] = OtherToThisCastSpec;
         casts[6] = NULL;

stringdtype/stringdtype/src/casts.h

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

stringdtype/stringdtype/src/dtype.c

@@ -533,9 +533,8 @@ init_string_dtype(void)
         PANDAS_AVAILABLE = 1;
     }
 
-    PyArrayMethod_Spec **StringDType_casts = get_casts(
-            (PyArray_DTypeMeta *)&StringDType,
-            (PyArray_DTypeMeta *)&PandasStringDType, PANDAS_AVAILABLE);
+    PyArrayMethod_Spec **StringDType_casts =
+            get_casts((PyArray_DTypeMeta *)&StringDType, NULL);
 
     PyArrayDTypeMeta_Spec StringDType_DTypeSpec = {
             .typeobj = StringScalar_Type,
@@ -557,10 +556,6 @@ 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;
@@ -585,7 +580,7 @@ init_string_dtype(void)
     if (PANDAS_AVAILABLE) {
         PyArrayMethod_Spec **PandasStringDType_casts =
                 get_casts((PyArray_DTypeMeta *)&PandasStringDType,
-                          (PyArray_DTypeMeta *)&StringDType, PANDAS_AVAILABLE);
+                          (PyArray_DTypeMeta *)&StringDType);
 
         PyArrayDTypeMeta_Spec PandasStringDType_DTypeSpec = {
                 .typeobj = PandasStringScalar_Type,
@@ -601,6 +596,8 @@ init_string_dtype(void)
             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);

stringdtype/tests/test_stringdtype.py

@@ -31,25 +31,18 @@ def dtype(request):
     if request.param == "StringDType":
         return StringDType()
     elif request.param == "PandasStringDType":
-        try:
-            from stringdtype import PandasStringDType
+        pytest.importorskip("pandas")
+        from stringdtype import PandasStringDType
 
-            return PandasStringDType()
-        except ImportError:
-            pytest.skip("cannot import pandas")
+        return PandasStringDType()
 
 
 @pytest.fixture
 def scalar(dtype):
     if dtype == StringDType():
         return StringScalar
-    try:
-        from stringdtype import PandasStringDType
-
-        del PandasStringDType
+    else:
         return PandasStringScalar
-    except ImportError:
-        pytest.skip("cannot import pandas")
 
 
 def test_scalar_creation(scalar):
@@ -77,6 +70,8 @@ def test_array_creation_utf8(dtype, data):
 
 
 def test_array_creation_scalars(string_list, scalar, dtype):
+    if not issubclass(scalar, dtype.type):
+        pytest.skip()
     arr = np.array([scalar(s) for s in string_list])
     assert repr(arr) == repr(np.array(string_list, dtype=dtype))
 
@@ -383,3 +378,25 @@ def test_create_with_na(dtype, na_val):
         == f"array(['hello', {dtype.na_object}, 'world'], dtype={dtype})"
     )
     assert arr[1] is dtype.na_object
+
+
+def test_pandas_to_numpy_cast(string_list):
+    pytest.importorskip("pandas")
+
+    from stringdtype import PandasStringDType
+
+    sarr = np.array(string_list, dtype=StringDType())
+
+    parr = sarr.astype(PandasStringDType())
+
+    np.testing.assert_array_equal(
+        parr, np.array(string_list, dtype=PandasStringDType())
+    )
+    np.testing.assert_array_equal(sarr, parr.astype(StringDType()))
+
+    # check that NA converts correctly too
+    sarr[1] = StringDType.na_object
+    parr = sarr.astype(PandasStringDType())
+
+    assert parr[1] is PandasStringDType.na_object
+    assert parr.astype(StringDType())[1] is StringDType.na_object