add support for pickling StringDType instances

pickling arrays is currently broken and needs support in numpy
to work

Author: ngoldbaum

stringdtype/stringdtype/__init__.py

@@ -1,7 +1,5 @@
-"""A dtype for working with string data
+"""A dtype for working with variable-length string data
 
-This is an example usage of the experimental new dtype API
-in Numpy and is not intended for any real purpose.
 """
 
 from .scalar import StringScalar  # isort: skip

stringdtype/stringdtype/src/dtype.c

@@ -200,6 +200,82 @@ stringdtype_repr(StringDTypeObject *NPY_UNUSED(self))
     return PyUnicode_FromString("StringDType()");
 }
 
+static int PICKLE_VERSION = 1;
+
+static PyObject *
+stringdtype__reduce__(StringDTypeObject *self)
+{
+    PyObject *ret, *mod, *obj, *state;
+
+    ret = PyTuple_New(3);
+    if (ret == NULL) {
+        return NULL;
+    }
+
+    mod = PyImport_ImportModule("stringdtype");
+    if (mod == NULL) {
+        Py_DECREF(ret);
+        return NULL;
+    }
+
+    obj = PyObject_GetAttrString(mod, "_reconstruct_StringDType");
+    Py_DECREF(mod);
+    if (obj == NULL) {
+        Py_DECREF(ret);
+        return NULL;
+    }
+
+    PyTuple_SET_ITEM(ret, 0, obj);
+
+    PyTuple_SET_ITEM(ret, 1, PyTuple_New(0));
+
+    state = PyTuple_New(1);
+
+    PyTuple_SET_ITEM(state, 0, PyLong_FromLong(PICKLE_VERSION));
+
+    PyTuple_SET_ITEM(ret, 2, state);
+
+    return ret;
+}
+
+static PyObject *
+stringdtype__setstate__(StringDTypeObject *NPY_UNUSED(self), PyObject *args)
+{
+    if (PyTuple_GET_SIZE(args) != 1 ||
+        !(PyLong_Check(PyTuple_GET_ITEM(args, 0)))) {
+        PyErr_BadInternalCall();
+        return NULL;
+    }
+
+    long version = PyLong_AsLong(PyTuple_GET_ITEM(args, 0));
+
+    if (version != PICKLE_VERSION) {
+        PyErr_Format(PyExc_ValueError,
+                     "Pickle version mismatch. Got version %d but expected "
+                     "version %d.",
+                     version, PICKLE_VERSION);
+        return NULL;
+    }
+
+    Py_RETURN_NONE;
+}
+
+static PyMethodDef StringDType_methods[] = {
+        {
+                "__reduce__",
+                (PyCFunction)stringdtype__reduce__,
+                METH_NOARGS,
+                "Reduction method for an StringDType object",
+        },
+        {
+                "__setstate__",
+                (PyCFunction)stringdtype__setstate__,
+                METH_O,
+                "Unpickle an StringDType object",
+        },
+        {NULL},
+};
+
 /*
  * This is the basic things that you need to create a Python Type/Class in C.
  * However, there is a slight difference here because we create a
@@ -215,6 +291,7 @@ PyArray_DTypeMeta StringDType = {
                 .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 */
 };

stringdtype/tests/test_stringdtype.py

@@ -1,3 +1,7 @@
+import os
+import pickle
+import tempfile
+
 import numpy as np
 import pytest
 
@@ -124,3 +128,17 @@ def test_memory_usage(string_list):
         _memory_usage("hello")
     with pytest.raises(TypeError):
         _memory_usage(np.array([1, 2, 3]))
+
+
+def test_pickle_dtype():
+    dtype = StringDType()
+
+    with tempfile.NamedTemporaryFile("wb", delete=False) as f:
+        pickle.dump(dtype, f)
+
+    with open(f.name, "rb") as f:
+        load_dtype = pickle.load(f)
+
+    assert dtype == load_dtype
+
+    os.remove(f.name)