Merge pull request #10860 from pv/pypyfix-npy-title-key

BUG: core: fix NPY_TITLE_KEY macro on pypy

2 files changed, 39 insertions, 2 deletions

diff --git a/numpy/core/include/numpy/ndarrayobject.h b/numpy/core/include/numpy/ndarrayobject.h
index 4e63868f3..ec0fd1ee9 100644
--- a/numpy/core/include/numpy/ndarrayobject.h
+++ b/numpy/core/include/numpy/ndarrayobject.h
@@ -232,9 +232,36 @@ PyArray_DiscardWritebackIfCopy(PyArrayObject *arr)
    dict.
 */
 
-#define NPY_TITLE_KEY(key, value) ((PyTuple_GET_SIZE((value))==3) && \
-                                   (PyTuple_GET_ITEM((value), 2) == (key)))
+static NPY_INLINE int
+NPY_TITLE_KEY_check(PyObject *key, PyObject *value)
+{
+    PyObject *title;
+    if (PyTuple_GET_SIZE(value) != 3) {
+        return 0;
+    }
+    title = PyTuple_GET_ITEM(value, 2);
+    if (key == title) {
+        return 1;
+    }
+#ifdef PYPY_VERSION
+    /*
+     * On PyPy, dictionary keys do not always preserve object identity.
+     * Fall back to comparison by value.
+     */
+    if (PyUnicode_Check(title) && PyUnicode_Check(key)) {
+        return PyUnicode_Compare(title, key) == 0 ? 1 : 0;
+    }
+#if PY_VERSION_HEX < 0x03000000
+    if (PyString_Check(title) && PyString_Check(key)) {
+        return PyObject_Compare(title, key) == 0 ? 1 : 0;
+    }
+#endif
+#endif
+    return 0;
+}
 
+/* Macro, for backward compat with "if NPY_TITLE_KEY(key, value) { ..." */
+#define NPY_TITLE_KEY(key, value) (NPY_TITLE_KEY_check((key), (value)))
 
 #define DEPRECATE(msg) PyErr_WarnEx(PyExc_DeprecationWarning,msg,1)
 #define DEPRECATE_FUTUREWARNING(msg) PyErr_WarnEx(PyExc_FutureWarning,msg,1)
diff --git a/numpy/core/tests/test_regression.py b/numpy/core/tests/test_regression.py
index fe0adb8cb..b3cb3e610 100644
--- a/numpy/core/tests/test_regression.py
+++ b/numpy/core/tests/test_regression.py
@@ -2332,3 +2332,13 @@ class TestRegression(object):
         #va[0] = b'\xff\xff\xff\xff'
         #del va
         #assert_equal(x, b'\x00\x00\x00\x00')
+
+    def test_structarray_title(self):
+        # The following used to segfault on pypy, due to NPY_TITLE_KEY
+        # not working properly and resulting to double-decref of the
+        # structured array field items:
+        # See: https://bitbucket.org/pypy/pypy/issues/2789
+        for j in range(5):
+            structure = np.array([1], dtype=[(('x', 'X'), np.object_)])
+            structure[0]['x'] = np.array([2])
+            gc.collect()