Add static PandasStringDType that uses pandas NA object

Author: ngoldbaum

stringdtype/stringdtype/__init__.py

@@ -3,12 +3,18 @@
 """
 
 from .missing import NA  # isort: skip
-from .scalar import StringScalar  # isort: skip
-from ._main import StringDType, _memory_usage
+from .scalar import StringScalar, PandasStringScalar  # isort: skip
+from ._main import PandasStringDType, StringDType, _memory_usage
 
 __all__ = [
     "NA",
     "StringDType",
     "StringScalar",
     "_memory_usage",
 ]
+
+# this happens when pandas isn't importable
+if StringDType is PandasStringDType:
+    del PandasStringDType
+else:
+    __all__.extend("PandasStringDType")

stringdtype/stringdtype/scalar.py

@@ -1,5 +1,9 @@
-"""A scalar type needed by the dtype machinery."""
+"""Scalar types needed by the dtype machinery."""
 
 
 class StringScalar(str):
     pass
+
+
+class PandasStringScalar(str):
+    pass

stringdtype/stringdtype/src/dtype.c

@@ -4,6 +4,7 @@
 #include "static_string.h"
 
 PyTypeObject *StringScalar_Type = NULL;
+PyTypeObject *PandasStringScalar_Type = NULL;
 static PyTypeObject *StringNA_Type = NULL;
 PyObject *NA_OBJ = NULL;
 
@@ -488,6 +489,30 @@ PyArray_DTypeMeta StringDType = {
         /* rest, filled in during DTypeMeta initialization */
 };
 
+/*
+ * Ideally we don't need the copy/pasted type below and we allow the
+ * StringDType class to be initialized dynamically however, that requires that
+ * the class is a heap type, and that likely only works well with some changes
+ * in numpy in Python 3.12 or newer. In practice we really only need a version
+ * with a numpy-specific NA and a and a version that uses Pandas' NA object,
+ * so we just define PandasStringDType statically as well and live with a bit
+ * of copy/paste
+ */
+
+PyArray_DTypeMeta PandasStringDType = {
+        {{
+                PyVarObject_HEAD_INIT(NULL, 0).tp_name =
+                        "stringdtype.PandasStringDType",
+                .tp_basicsize = sizeof(StringDTypeObject),
+                .tp_new = stringdtype_new,
+                .tp_dealloc = (destructor)stringdtype_dealloc,
+                .tp_repr = (reprfunc)stringdtype_repr,
+                .tp_str = (reprfunc)stringdtype_repr,
+                .tp_methods = StringDType_methods,
+        }},
+        /* rest, filled in during DTypeMeta initialization */
+};
+
 int
 init_string_dtype(void)
 {
@@ -522,6 +547,51 @@ init_string_dtype(void)
 
     StringDType.singleton = singleton;
 
+    /* and once again for PandasStringDType */
+
+    PyObject *mod = PyImport_ImportModule("pandas");
+
+    if (mod != NULL) {
+        PyArrayDTypeMeta_Spec PandasStringDType_DTypeSpec = {
+                .typeobj = PandasStringScalar_Type,
+                .slots = StringDType_Slots,
+                .casts = casts,
+        };
+
+        PyObject *pandas_na_obj = PyObject_GetAttrString(mod, "NA");
+
+        Py_DECREF(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();
+        PyDict_SetItemString(((PyTypeObject *)&PandasStringDType)->tp_dict,
+                             "na_object", pandas_na_obj);
+        if (PyType_Ready((PyTypeObject *)&PandasStringDType) < 0) {
+            return -1;
+        }
+
+        if (PyArrayInitDTypeMeta_FromSpec(&PandasStringDType,
+                                          &PandasStringDType_DTypeSpec) < 0) {
+            return -1;
+        }
+
+        singleton = PyArray_GetDefaultDescr(&PandasStringDType);
+
+        if (singleton == NULL) {
+            return -1;
+        }
+
+        PandasStringDType.singleton = singleton;
+    }
+    else {
+        PandasStringDType = StringDType;
+    }
+
     for (int i = 0; casts[i] != NULL; i++) {
         free(casts[i]->dtypes);
         free(casts[i]);

stringdtype/stringdtype/src/dtype.h

@@ -19,7 +19,9 @@ typedef struct {
 } StringDTypeObject;
 
 extern PyArray_DTypeMeta StringDType;
+extern PyArray_DTypeMeta PandasStringDType;
 extern PyTypeObject *StringScalar_Type;
+extern PyTypeObject *PandasStringScalar_Type;
 extern PyObject *NA_OBJ;
 
 PyObject *

stringdtype/stringdtype/src/main.c

@@ -107,12 +107,20 @@ PyInit__main(void)
 
     StringScalar_Type =
             (PyTypeObject *)PyObject_GetAttrString(mod, "StringScalar");
-    Py_DECREF(mod);
 
     if (StringScalar_Type == NULL) {
         goto error;
     }
 
+    PandasStringScalar_Type =
+            (PyTypeObject *)PyObject_GetAttrString(mod, "PandasStringScalar");
+
+    Py_DECREF(mod);
+
+    if (PandasStringScalar_Type == NULL) {
+        goto error;
+    }
+
     if (init_string_na_object(mod) < 0) {
         goto error;
     }
@@ -127,6 +135,13 @@ PyInit__main(void)
         goto error;
     }
 
+    Py_INCREF((PyObject *)&PandasStringDType);
+    if (PyModule_AddObject(m, "PandasStringDType",
+                           (PyObject *)&PandasStringDType) < 0) {
+        Py_DECREF((PyObject *)&PandasStringDType);
+        goto error;
+    }
+
     if (init_ufuncs() < 0) {
         goto error;
     }

stringdtype/tests/test_stringdtype.py

@@ -338,3 +338,15 @@ def test_create_with_na(na_val):
         == "array(['hello', stringdtype.NA, 'world'], dtype=StringDType())"
     )
     assert arr[1] == NA and arr[1] is NA
+
+
+def test_pandas_string_dtype():
+    pandas = pytest.importorskip("pandas")
+    from stringdtype import PandasStringDType
+
+    assert PandasStringDType.na_object is pandas.NA
+
+    string_list = ["hello", np.nan, "world"]
+    arr = np.array(string_list, dtype=PandasStringDType())
+
+    assert arr[1] is pandas.NA