BUG: Fix dtype refcount in __array__ (#29715)

mtsokol · web-flow · commit c9e75927a43c · 2025-09-17T08:35:15.000+02:00
* BUG: Fix `dtype` refcount in `__array__`

* Consider all possible code paths

* Remove else-if branch

* Move refcount checks to a separate test

* Add code comments

* Add missing `Py_DECREF` for error path

* Apply review comments
diff --git a/numpy/_core/src/multiarray/methods.c b/numpy/_core/src/multiarray/methods.c
@@ -934,6 +934,11 @@ array_getarray(PyArrayObject *self, PyObject *args, PyObject *kwds)
         return NULL;
     }
 
+    if (newtype == NULL) {
+        newtype = PyArray_DESCR(self);
+        Py_INCREF(newtype);  // newtype is owned.
+    }
+
     /* convert to PyArray_Type */
     if (!PyArray_CheckExact(self)) {
         PyArrayObject *new;
@@ -951,6 +956,7 @@ array_getarray(PyArrayObject *self, PyObject *args, PyObject *kwds)
                 (PyObject *)self
         );
         if (new == NULL) {
+            Py_DECREF(newtype);
             return NULL;
         }
         self = new;
@@ -960,22 +966,21 @@ array_getarray(PyArrayObject *self, PyObject *args, PyObject *kwds)
     }
 
     if (copy == NPY_COPY_ALWAYS) {
-        if (newtype == NULL) {
-            newtype = PyArray_DESCR(self);
-        }
-        ret = PyArray_CastToType(self, newtype, 0);
+        ret = PyArray_CastToType(self, newtype, 0);  // steals newtype reference
         Py_DECREF(self);
         return ret;
     } else { // copy == NPY_COPY_IF_NEEDED || copy == NPY_COPY_NEVER
-        if (newtype == NULL || PyArray_EquivTypes(PyArray_DESCR(self), newtype)) {
+        if (PyArray_EquivTypes(PyArray_DESCR(self), newtype)) {
+            Py_DECREF(newtype);
             return (PyObject *)self;
         }
         if (copy == NPY_COPY_IF_NEEDED) {
-            ret = PyArray_CastToType(self, newtype, 0);
+            ret = PyArray_CastToType(self, newtype, 0);  // steals newtype reference.
             Py_DECREF(self);
             return ret;
         } else { // copy == NPY_COPY_NEVER
             PyErr_SetString(PyExc_ValueError, npy_no_copy_err_msg);
+            Py_DECREF(newtype);
             Py_DECREF(self);
             return NULL;
         }
diff --git a/numpy/_core/tests/test_api.py b/numpy/_core/tests/test_api.py
@@ -91,7 +91,7 @@ def test_array_array():
     # instead we get a array([...], dtype=">V18")
     assert_equal(bytes(np.array(o).data), bytes(a.data))
 
-    # test array
+    # test __array__
     def custom__array__(self, dtype=None, copy=None):
         return np.array(100.0, dtype=dtype, copy=copy)
 
@@ -157,6 +157,39 @@ def custom__array__(self, dtype=None, copy=None):
     assert_equal(np.array([(1.0,) * 10] * 10, dtype=np.float64),
                  np.ones((10, 10), dtype=np.float64))
 
+
+@pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
+def test___array___refcount():
+    class MyArray:
+        def __init__(self, dtype):
+            self.val = np.array(-1, dtype=dtype)
+
+        def __array__(self, dtype=None, copy=None):
+            return self.val.__array__(dtype=dtype, copy=copy)
+
+    # test all possible scenarios:
+    # dtype(none | same | different) x copy(true | false | none)
+    dt = np.dtype(np.int32)
+    old_refcount = sys.getrefcount(dt)
+    np.array(MyArray(dt))
+    assert_equal(old_refcount, sys.getrefcount(dt))
+    np.array(MyArray(dt), dtype=dt)
+    assert_equal(old_refcount, sys.getrefcount(dt))
+    np.array(MyArray(dt), copy=None)
+    assert_equal(old_refcount, sys.getrefcount(dt))
+    np.array(MyArray(dt), dtype=dt, copy=None)
+    assert_equal(old_refcount, sys.getrefcount(dt))
+    dt2 = np.dtype(np.int16)
+    old_refcount2 = sys.getrefcount(dt2)
+    np.array(MyArray(dt), dtype=dt2)
+    assert_equal(old_refcount2, sys.getrefcount(dt2))
+    np.array(MyArray(dt), dtype=dt2, copy=None)
+    assert_equal(old_refcount2, sys.getrefcount(dt2))
+    with pytest.raises(ValueError):
+        np.array(MyArray(dt), dtype=dt2, copy=False)
+    assert_equal(old_refcount2, sys.getrefcount(dt2))
+
+
 @pytest.mark.parametrize("array", [True, False])
 def test_array_impossible_casts(array):
     # All builtin types can be forcibly cast, at least theoretically,
