MAINT: return early from PyArray_FillObjectArray for non-legacy dtypes

2 files changed, 17 insertions, 24 deletions

diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
index a791dc35c..38af60427 100644
--- a/numpy/core/src/multiarray/ctors.c
+++ b/numpy/core/src/multiarray/ctors.c
@@ -1076,27 +1076,11 @@ PyArray_NewLikeArrayWithShape(PyArrayObject *prototype, NPY_ORDER order,
     }
 
     /* Logic shared by `empty`, `empty_like`, and `ndarray.__new__` */
-    PyArray_Descr* descr = PyArray_DESCR((PyArrayObject *)ret);
-    if (PyDataType_REFCHK(descr)) {
-        if (NPY_DT_is_legacy(NPY_DTYPE(dtype))) {
-            PyArray_FillObjectArray((PyArrayObject *)ret, Py_None);
-            if (PyErr_Occurred()) {
-                Py_DECREF(ret);
-                return NULL;
-            }
-        }
-        else {
-            // Currently we assume all non-legacy dtypes with the
-            // NPY_ITEM_REFCOUNT flag either hold heap-allocated data or hold
-            // python objects. In the latter case the dtype is responsible for
-            // managing initialization and reference counts. In both cases
-            // initializing to NULL makes sense.
-            char *optr = PyArray_DATA((PyArrayObject*)ret);
-            npy_intp n = PyArray_SIZE((PyArrayObject*)ret);
-            for (npy_intp i = 0; i < n; i++) {
-                optr = NULL;
-                optr += descr->elsize;
-            }
+    if (PyDataType_REFCHK(PyArray_DESCR((PyArrayObject *)ret))) {
+        PyArray_FillObjectArray((PyArrayObject *)ret, Py_None);
+        if (PyErr_Occurred()) {
+            Py_DECREF(ret);
+            return NULL;
         }
     }
 
diff --git a/numpy/core/src/multiarray/refcount.c b/numpy/core/src/multiarray/refcount.c
index e735e3b8f..20527f7af 100644
--- a/numpy/core/src/multiarray/refcount.c
+++ b/numpy/core/src/multiarray/refcount.c
@@ -16,6 +16,7 @@
 #include "numpy/arrayobject.h"
 #include "numpy/arrayscalars.h"
 #include "iterators.h"
+#include "dtypemeta.h"
 
 #include "npy_config.h"
 
@@ -358,9 +359,17 @@ PyArray_XDECREF(PyArrayObject *mp)
 NPY_NO_EXPORT void
 PyArray_FillObjectArray(PyArrayObject *arr, PyObject *obj)
 {
+    PyArray_Descr* descr = PyArray_DESCR(arr);
+
+    // non-legacy dtypes are responsible for initializing
+    // their own internal references
+    if (!NPY_DT_is_legacy(NPY_DTYPE(descr))) {
+        return;
+    }
+
     npy_intp i,n;
     n = PyArray_SIZE(arr);
-    if (PyArray_DESCR(arr)->type_num == NPY_OBJECT) {
+    if (descr->type_num == NPY_OBJECT) {
         PyObject **optr;
         optr = (PyObject **)(PyArray_DATA(arr));
         n = PyArray_SIZE(arr);
@@ -380,8 +389,8 @@ PyArray_FillObjectArray(PyArrayObject *arr, PyObject *obj)
         char *optr;
         optr = PyArray_DATA(arr);
         for (i = 0; i < n; i++) {
-            _fillobject(optr, obj, PyArray_DESCR(arr));
-            optr += PyArray_DESCR(arr)->elsize;
+            _fillobject(optr, obj, descr);
+            optr += descr->elsize;
         }
     }
 }