MAINT: some hacks to fix pypy support

Author: ngoldbaum

numpy/_core/src/multiarray/stringdtype/dtype.c

@@ -248,6 +248,27 @@ string_discover_descriptor_from_pyobject(PyTypeObject *NPY_UNUSED(cls),
     return ret;
 }
 
+#ifndef PYPY_VERSION
+#define NA_IS_COMPARE(a, b) a == b
+#else
+static PyObject *py_id = NULL;
+
+static int
+eq_compare_pypy_na(PyObject *a, PyObject *b) {
+    PyObject *a_id = PyObject_CallOneArg(py_id, a);
+    if (a_id == NULL) {
+        return -1;
+    }
+    PyObject *b_id = PyObject_CallOneArg(py_id, b);
+    if (b_id == NULL) {
+        return -1;
+    }
+    return PyObject_RichCompareBool(a_id, b_id, Py_EQ);
+}
+
+#define NA_IS_COMPARE(a, b) eq_compare_pypy_na(a, b)
+#endif
+
 // Take a python object `obj` and insert it into the array of dtype `descr` at
 // the position given by dataptr.
 int
@@ -260,40 +281,47 @@ stringdtype_setitem(StringDTypeObject *descr, PyObject *obj, char **dataptr)
     // borrow reference
     PyObject *na_object = descr->na_object;
 
-    // setting NA *must* check pointer equality since NA types might not
-    // allow equality
-    if (na_object != NULL && obj == na_object) {
-        if (NpyString_pack_null(allocator, sdata) < 0) {
-            PyErr_SetString(PyExc_MemoryError,
-                            "Failed to pack null string during StringDType "
-                            "setitem");
+    if (na_object != NULL) {
+        // specifically for pandas.NA, we *have* to use "is" comparison
+        // because there's no other way to identify NA, since "NA == anything"
+        // returns NA.
+        int is_cmp = NA_IS_COMPARE(obj, na_object);
+        if (is_cmp == -1) {
             goto fail;
         }
+        if (is_cmp) {
+            if (NpyString_pack_null(allocator, sdata) < 0) {
+                PyErr_SetString(PyExc_MemoryError,
+                                "Failed to pack null string during StringDType "
+                                "setitem");
+                goto fail;
+            }
+            goto success;
+        }
     }
-    else {
-        PyObject *val_obj = get_value(obj, descr->coerce);
+    PyObject *val_obj = get_value(obj, descr->coerce);
 
-        if (val_obj == NULL) {
-            goto fail;
-        }
+    if (val_obj == NULL) {
+        goto fail;
+    }
 
-        Py_ssize_t length = 0;
-        const char *val = PyUnicode_AsUTF8AndSize(val_obj, &length);
-        if (val == NULL) {
-            Py_DECREF(val_obj);
-            goto fail;
-        }
+    Py_ssize_t length = 0;
+    const char *val = PyUnicode_AsUTF8AndSize(val_obj, &length);
+    if (val == NULL) {
+        Py_DECREF(val_obj);
+        goto fail;
+    }
 
-        if (NpyString_pack(allocator, sdata, val, length) < 0) {
-            PyErr_SetString(PyExc_MemoryError,
-                            "Failed to pack string during StringDType "
-                            "setitem");
-            Py_DECREF(val_obj);
-            goto fail;
-        }
+    if (NpyString_pack(allocator, sdata, val, length) < 0) {
+        PyErr_SetString(PyExc_MemoryError,
+                        "Failed to pack string during StringDType "
+                        "setitem");
         Py_DECREF(val_obj);
+        goto fail;
     }
+    Py_DECREF(val_obj);
 
+success:
     NpyString_release_allocator(descr);
 
     return 0;
@@ -853,6 +881,13 @@ init_string_dtype(void)
             .casts = StringDType_casts,
     };
 
+#ifdef PYPY_VERSION
+    py_id = PyDict_GetItemString(PyEval_GetBuiltins(), "id");
+    if (py_id == NULL) {
+        return -1;
+    }
+#endif
+
     /* Loaded dynamically, so may need to be set here: */
     ((PyObject *)&StringDType)->ob_type = &PyArrayDTypeMeta_Type;
     ((PyTypeObject *)&StringDType)->tp_base = &PyArrayDescr_Type;