Merge pull request #68 from ngoldbaum/stringdtype-view-casts

Add casts between StringDType and PandasStringDType

Author: ngoldbaum

stringdtype/pyproject.toml

@@ -35,6 +35,6 @@ per-file-ignores = {"__init__.py" = ["F401"]}
 
 [tool.meson-python.args]
 dist = []
-setup = ["-Ddebug=true", "-Doptimization=0"]
+setup = ["-Dbuildtype=debug"]
 compile = []
 install = []

stringdtype/stringdtype/src/casts.c

@@ -75,6 +75,9 @@ 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";
 
 // unicode to string
 
@@ -476,38 +479,80 @@ get_dtypes(PyArray_DTypeMeta *dt1, PyArray_DTypeMeta *dt2)
 }
 
 PyArrayMethod_Spec **
-get_casts(void)
+get_casts(PyArray_DTypeMeta *this, PyArray_DTypeMeta *other)
 {
-    PyArray_DTypeMeta **s2s_dtypes = get_dtypes(NULL, NULL);
+    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;
+
+    int is_pandas = (this == (PyArray_DTypeMeta *)&PandasStringDType);
+
+    int num_casts = 5;
+
+    if (is_pandas) {
+        num_casts = 7;
+
+        PyArray_DTypeMeta **t2o_dtypes = get_dtypes(this, other);
 
-    PyArray_DTypeMeta **u2s_dtypes = get_dtypes(&PyArray_UnicodeDType, NULL);
+        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(s2p_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 *));
-    casts[0] = StringToStringCastSpec;
+    PyArrayMethod_Spec **casts = NULL;
+
+    casts = malloc(num_casts * sizeof(PyArrayMethod_Spec *));
+
+    casts[0] = ThisToThisCastSpec;
     casts[1] = UnicodeToStringCastSpec;
     casts[2] = StringToUnicodeCastSpec;
     casts[3] = StringToBoolCastSpec;
-    casts[4] = NULL;
+    if (is_pandas) {
+        casts[4] = ThisToOtherCastSpec;
+        casts[5] = OtherToThisCastSpec;
+        casts[6] = NULL;
+    }
+    else {
+        casts[4] = NULL;
+    }
 
     return casts;
 }

stringdtype/stringdtype/src/casts.h

@@ -11,7 +11,7 @@
 #include "numpy/ndarraytypes.h"
 
 PyArrayMethod_Spec **
-get_casts(void);
+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

@@ -523,12 +523,23 @@ 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, NULL);
 
     PyArrayDTypeMeta_Spec StringDType_DTypeSpec = {
             .typeobj = StringScalar_Type,
             .slots = StringDType_Slots,
-            .casts = casts,
+            .casts = StringDType_casts,
     };
 
     /* Loaded dynamically, so may need to be set here: */
@@ -559,20 +570,27 @@ 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);
 
-    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;
@@ -605,15 +623,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;

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