Merge pull request #81 from ngoldbaum/no-coerce

Optionally do not coerce non-strings to string in setitem

Author: ngoldbaum

stringdtype/stringdtype/missing.py

@@ -1,4 +1,13 @@
-class NAType:
+class Singleton:
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+
+class NAType(Singleton):
     def __repr__(self):
         return "stringdtype.NA"
 

stringdtype/stringdtype/src/casts.c

@@ -21,7 +21,7 @@ gil_error(PyObject *type, const char *msg)
     {                                                                      \
         if (given_descrs[1] == NULL) {                                     \
             PyArray_Descr *new =                                           \
-                    (PyArray_Descr *)new_stringdtype_instance(NA_OBJ);     \
+                    (PyArray_Descr *)new_stringdtype_instance(NA_OBJ, 0);  \
             if (new == NULL) {                                             \
                 return (NPY_CASTING)-1;                                    \
             }                                                              \

stringdtype/stringdtype/src/dtype.c

@@ -11,7 +11,7 @@ PyObject *NA_OBJ = NULL;
  * Internal helper to create new instances
  */
 PyObject *
-new_stringdtype_instance(PyObject *na_object)
+new_stringdtype_instance(PyObject *na_object, int coerce)
 {
     PyObject *new =
             PyArrayDescr_Type.tp_new((PyTypeObject *)&StringDType, NULL, NULL);
@@ -22,6 +22,7 @@ new_stringdtype_instance(PyObject *na_object)
 
     Py_INCREF(na_object);
     ((StringDTypeObject *)new)->na_object = na_object;
+    ((StringDTypeObject *)new)->coerce = coerce;
 
     PyArray_Descr *base = (PyArray_Descr *)new;
     base->elsize = sizeof(ss);
@@ -67,23 +68,32 @@ common_dtype(PyArray_DTypeMeta *cls, PyArray_DTypeMeta *other)
 }
 
 // returns a new reference to the string "value" of
-// `scalar`. If scalar is not already a string, __str__
-// is called to convert it to a string. If the scalar
-// is the na_object for the dtype class, return
-// a new reference to the na_object.
+// `scalar`. If scalar is not already a string and
+// coerce is nonzero, __str__ is called to convert it
+// to a string. If coerce is zero, raises an error for
+// non-string or non-NA input. If the scalar is the
+// na_object for the dtype class, return a new
+// reference to the na_object.
 
 static PyObject *
-get_value(PyObject *scalar)
+get_value(PyObject *scalar, int coerce)
 {
     PyTypeObject *scalar_type = Py_TYPE(scalar);
     if (!((scalar_type == &PyUnicode_Type) ||
           (scalar_type == StringScalar_Type))) {
-        // attempt to coerce to str
-        scalar = PyObject_Str(scalar);
-        if (scalar == NULL) {
-            // __str__ raised an exception
+        if (coerce == 0) {
+            PyErr_SetString(PyExc_ValueError,
+                            "StringDType only allows string data");
             return NULL;
         }
+        else {
+            // attempt to coerce to str
+            scalar = PyObject_Str(scalar);
+            if (scalar == NULL) {
+                // __str__ raised an exception
+                return NULL;
+            }
+        }
     }
     // attempt to decode as UTF8
     return PyUnicode_AsUTF8String(scalar);
@@ -93,12 +103,12 @@ static PyArray_Descr *
 string_discover_descriptor_from_pyobject(PyTypeObject *NPY_UNUSED(cls),
                                          PyObject *obj)
 {
-    PyObject *val = get_value(obj);
+    PyObject *val = get_value(obj, 1);
     if (val == NULL) {
         return NULL;
     }
 
-    PyArray_Descr *ret = (PyArray_Descr *)new_stringdtype_instance(NA_OBJ);
+    PyArray_Descr *ret = (PyArray_Descr *)new_stringdtype_instance(NA_OBJ, 1);
     if (ret == NULL) {
         return NULL;
     }
@@ -126,7 +136,7 @@ stringdtype_setitem(StringDTypeObject *descr, PyObject *obj, char **dataptr)
         // so it already contains a NA value
     }
     else {
-        PyObject *val_obj = get_value(obj);
+        PyObject *val_obj = get_value(obj, descr->coerce);
 
         if (val_obj == NULL) {
             return -1;
@@ -334,21 +344,23 @@ static PyType_Slot StringDType_Slots[] = {
 static PyObject *
 stringdtype_new(PyTypeObject *NPY_UNUSED(cls), PyObject *args, PyObject *kwds)
 {
-    static char *kwargs_strs[] = {"size", "na_object", NULL};
+    static char *kwargs_strs[] = {"size", "na_object", "coerce", NULL};
 
     long size = 0;
     PyObject *na_object = NULL;
+    int coerce = 1;
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|lO:StringDType",
-                                     kwargs_strs, &size, &na_object)) {
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|lOp:StringDType",
+                                     kwargs_strs, &size, &na_object,
+                                     &coerce)) {
         return NULL;
     }
 
     if (na_object == NULL) {
         na_object = NA_OBJ;
     }
 
-    PyObject *ret = new_stringdtype_instance(na_object);
+    PyObject *ret = new_stringdtype_instance(na_object, coerce);
 
     return ret;
 }
@@ -365,11 +377,18 @@ stringdtype_repr(StringDTypeObject *self)
     PyObject *ret = NULL;
     // borrow reference
     PyObject *na_object = self->na_object;
+    int coerce = self->coerce;
 
     // TODO: handle non-default NA
-    if (na_object != NA_OBJ) {
-        ret = PyUnicode_FromFormat("StringDType(na_object=%R)",
-                                   self->na_object);
+    if (na_object != NA_OBJ && coerce == 0) {
+        ret = PyUnicode_FromFormat("StringDType(na_object=%R, coerce=False)",
+                                   na_object);
+    }
+    else if (na_object != NA_OBJ) {
+        ret = PyUnicode_FromFormat("StringDType(na_object=%R)", na_object);
+    }
+    else if (coerce == 0) {
+        ret = PyUnicode_FromFormat("StringDType(coerce=False)", coerce);
     }
     else {
         ret = PyUnicode_FromString("StringDType()");
@@ -378,7 +397,7 @@ stringdtype_repr(StringDTypeObject *self)
     return ret;
 }
 
-static int PICKLE_VERSION = 1;
+static int PICKLE_VERSION = 2;
 
 static PyObject *
 stringdtype__reduce__(StringDTypeObject *self)
@@ -405,9 +424,9 @@ stringdtype__reduce__(StringDTypeObject *self)
 
     PyTuple_SET_ITEM(ret, 0, obj);
 
-    PyTuple_SET_ITEM(
-            ret, 1,
-            Py_BuildValue("(NO)", PyLong_FromLong(0), self->na_object));
+    PyTuple_SET_ITEM(ret, 1,
+                     Py_BuildValue("(NOi)", PyLong_FromLong(0),
+                                   self->na_object, self->coerce));
 
     PyTuple_SET_ITEM(ret, 2, Py_BuildValue("(l)", PICKLE_VERSION));
 
@@ -456,9 +475,39 @@ static PyMemberDef StringDType_members[] = {
         {"na_object", T_OBJECT_EX, offsetof(StringDTypeObject, na_object),
          READONLY,
          "The missing value object associated with the dtype instance"},
+        {"coerce", T_INT, offsetof(StringDTypeObject, coerce), READONLY,
+         "Controls hether non-string values should be coerced to string"},
         {NULL, 0, 0, 0, NULL},
 };
 
+static PyObject *
+StringDType_richcompare(PyObject *self, PyObject *other, int op)
+{
+    if (!((op == Py_EQ) || (op == Py_NE)) ||
+        (Py_TYPE(other) != Py_TYPE(self))) {
+        Py_INCREF(Py_NotImplemented);
+        return Py_NotImplemented;
+    }
+
+    // we know both are instances of StringDType so this is safe
+    StringDTypeObject *sself = (StringDTypeObject *)self;
+    StringDTypeObject *sother = (StringDTypeObject *)other;
+
+    int eq = (sself->na_object == sother->na_object) &&
+             (sself->coerce == sother->coerce);
+
+    PyObject *ret = Py_NotImplemented;
+    if ((op == Py_EQ && eq) || (op == Py_NE && !eq)) {
+        ret = Py_True;
+    }
+    else {
+        ret = Py_False;
+    }
+
+    Py_INCREF(ret);
+    return ret;
+}
+
 /*
  * 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
@@ -476,6 +525,7 @@ StringDType_type StringDType = {
                 .tp_str = (reprfunc)stringdtype_repr,
                 .tp_methods = StringDType_methods,
                 .tp_members = StringDType_members,
+                .tp_richcompare = StringDType_richcompare,
         }}},
         /* rest, filled in during DTypeMeta initialization */
 };

stringdtype/stringdtype/src/dtype.h

@@ -20,21 +20,19 @@
 typedef struct {
     PyArray_Descr base;
     PyObject *na_object;
+    int coerce;
 } StringDTypeObject;
 
 typedef struct {
     PyArray_DTypeMeta base;
 } StringDType_type;
 
 extern StringDType_type StringDType;
-extern StringDType_type PandasStringDType;
 extern PyTypeObject *StringScalar_Type;
-extern PyTypeObject *PandasStringScalar_Type;
 extern PyObject *NA_OBJ;
-extern int PANDAS_AVAILABLE;
 
 PyObject *
-new_stringdtype_instance(PyObject *na_object);
+new_stringdtype_instance(PyObject *na_object, int coerce);
 
 int
 init_string_dtype(void);

stringdtype/stringdtype/src/umath.c

@@ -19,17 +19,17 @@ multiply_resolve_descriptors(
     Py_INCREF(rdescr);
     loop_descrs[1] = rdescr;
 
-    PyArray_Descr *odescr = NULL;
+    StringDTypeObject *odescr = NULL;
 
     if (dtypes[0] == (PyArray_DTypeMeta *)&StringDType) {
-        odescr = ldescr;
+        odescr = (StringDTypeObject *)ldescr;
     }
     else {
-        odescr = rdescr;
+        odescr = (StringDTypeObject *)rdescr;
     }
 
     loop_descrs[2] = (PyArray_Descr *)new_stringdtype_instance(
-            ((StringDTypeObject *)odescr)->na_object);
+            odescr->na_object, odescr->coerce);
 
     return NPY_NO_CASTING;
 }

stringdtype/tests/test_stringdtype.py

@@ -26,11 +26,35 @@ def string_list():
 )
 
 
+@pytest.fixture(params=[True, False])
+def coerce(request):
+    return request.param
+
+
 @pytest.fixture(
     params=[None, NA, pd_param], ids=["None", "stringdtype.NA", "pandas.NA"]
 )
-def dtype(request):
-    return StringDType(na_object=request.param)
+def na_object(request):
+    return request.param
+
+
+@pytest.fixture()
+def dtype(na_object, coerce):
+    return StringDType(na_object=na_object, coerce=coerce)
+
+
+def test_dtype_creation():
+    dt = StringDType()
+    assert dt.na_object is NA and dt.coerce == 1
+
+    dt = StringDType(na_object=None)
+    assert dt.na_object is None and dt.coerce == 1
+
+    dt = StringDType(coerce=False)
+    assert dt.na_object is NA and dt.coerce == 0
+
+    dt = StringDType(na_object=None, coerce=False)
+    assert dt.na_object is None and dt.coerce == 0
 
 
 def test_scalar_creation():
@@ -44,10 +68,17 @@ def test_dtype_equality(dtype):
 
 
 def test_dtype_repr(dtype):
-    if dtype.na_object is NA:
+    if dtype.na_object is NA and dtype.coerce == 1:
         assert repr(dtype) == "StringDType()"
-    else:
+    elif dtype.coerce == 1:
         assert repr(dtype) == f"StringDType(na_object={dtype.na_object})"
+    elif dtype.na_object is NA:
+        assert repr(dtype) == "StringDType(coerce=False)"
+    else:
+        assert (
+            repr(dtype)
+            == f"StringDType(na_object={dtype.na_object}, coerce=False)"
+        )
 
 
 @pytest.mark.parametrize(
@@ -61,12 +92,17 @@ def test_dtype_repr(dtype):
 )
 def test_array_creation_utf8(dtype, data):
     arr = np.array(data, dtype=dtype)
-    assert repr(arr) == f"array({str(data)}, dtype={dtype})"
+    assert str(arr) == "[" + " ".join(["'" + str(d) + "'" for d in data]) + "]"
+    assert arr.dtype == dtype
 
 
 def test_array_creation_scalars(string_list):
     arr = np.array([StringScalar(s) for s in string_list])
-    assert repr(arr) == repr(np.array(string_list, dtype=StringDType()))
+    assert (
+        str(arr)
+        == "[" + " ".join(["'" + str(s) + "'" for s in string_list]) + "]"
+    )
+    assert arr.dtype == StringDType()
 
 
 @pytest.mark.parametrize(
@@ -77,11 +113,15 @@ def test_array_creation_scalars(string_list):
         [object, object, object],
     ],
 )
-def test_scalars_string_conversion(data):
-    np.testing.assert_array_equal(
-        np.array(data, dtype=StringDType()),
-        np.array([str(d) for d in data], dtype=StringDType()),
-    )
+def test_scalars_string_conversion(data, dtype):
+    if dtype.coerce != 0:
+        np.testing.assert_array_equal(
+            np.array(data, dtype=dtype),
+            np.array([str(d) for d in data], dtype=dtype),
+        )
+    else:
+        with pytest.raises(ValueError):
+            np.array(data, dtype=dtype)
 
 
 @pytest.mark.parametrize(
@@ -516,10 +556,7 @@ def test_create_with_na(dtype):
     na_val = dtype.na_object
     string_list = ["hello", na_val, "world"]
     arr = np.array(string_list, dtype=dtype)
-    assert (
-        repr(arr)
-        == f"array(['hello', {dtype.na_object}, 'world'], dtype={dtype})"
-    )
+    assert str(arr) == "[" + " ".join([repr(s) for s in string_list]) + "]"
     assert arr[1] is dtype.na_object