ENH: Use array_subscript_asarray to ensure an array type return value

3 files changed, 12 insertions, 2 deletions

diff --git a/numpy/core/src/multiarray/arrayobject.c b/numpy/core/src/multiarray/arrayobject.c
index b559d4aba..ad196e010 100644
--- a/numpy/core/src/multiarray/arrayobject.c
+++ b/numpy/core/src/multiarray/arrayobject.c
@@ -1172,12 +1172,12 @@ _void_compare(PyArrayObject *self, PyArrayObject *other, int cmp_op)
             if NPY_TITLE_KEY(key, value) {
                 continue;
             }
-            a = PyArray_EnsureAnyArray(array_subscript(self, key));
+            a = array_subscript_asarray(self, key);
             if (a == NULL) {
                 Py_XDECREF(res);
                 return NULL;
             }
-            b = array_subscript(other, key);
+            b = array_subscript_asarray(other, key);
             if (b == NULL) {
                 Py_XDECREF(res);
                 Py_DECREF(a);
diff --git a/numpy/core/src/multiarray/mapping.c b/numpy/core/src/multiarray/mapping.c
index a4b1e6575..c3b33098c 100644
--- a/numpy/core/src/multiarray/mapping.c
+++ b/numpy/core/src/multiarray/mapping.c
@@ -994,6 +994,13 @@ array_subscript_fancy(PyArrayObject *self, PyObject *op, int fancy)
     return other;
 }
 
+/* make sure subscript always returns an array object */
+NPY_NO_EXPORT PyObject *
+array_subscript_asarray(PyArrayObject *self, PyObject *op)
+{
+    return PyArray_EnsureAnyArray(array_subscript(self, op));
+}
+
 NPY_NO_EXPORT PyObject *
 array_subscript(PyArrayObject *self, PyObject *op)
 {
diff --git a/numpy/core/src/multiarray/mapping.h b/numpy/core/src/multiarray/mapping.h
index 0f697c373..ef08ca4f1 100644
--- a/numpy/core/src/multiarray/mapping.h
+++ b/numpy/core/src/multiarray/mapping.h
@@ -20,6 +20,9 @@ NPY_NO_EXPORT PyObject *
 array_item(PyArrayObject *self, Py_ssize_t i);
 
 NPY_NO_EXPORT PyObject *
+array_subscript_asarray(PyArrayObject *self, PyObject *op);
+
+NPY_NO_EXPORT PyObject *
 array_subscript(PyArrayObject *self, PyObject *op);
 
 NPY_NO_EXPORT int