Merge pull request numpy#27343 from lysnikolaou/flatindex-on-flatiter

ngoldbaum · web-flow · commit 760dbe9e6524 · 2025-01-29T16:18:42.000-07:00
diff --git a/numpy/_core/src/multiarray/iterators.c b/numpy/_core/src/multiarray/iterators.c
@@ -694,9 +694,23 @@ iter_subscript(PyArrayIterObject *self, PyObject *ind)
         obj = ind;
     }
 
-    /* Any remaining valid input is an array or has been turned into one */
     if (!PyArray_Check(obj)) {
-        goto fail;
+        PyArrayObject *tmp_arr = (PyArrayObject *) PyArray_FROM_O(obj);
+        if (tmp_arr == NULL) {
+            goto fail;
+        }
+
+        if (PyArray_SIZE(tmp_arr) == 0) {
+            PyArray_Descr *indtype = PyArray_DescrFromType(NPY_INTP);
+            Py_SETREF(obj, PyArray_FromArray(tmp_arr, indtype, NPY_ARRAY_FORCECAST));
+            Py_DECREF(tmp_arr);
+            if (obj == NULL) {
+                goto fail;
+            }
+        }
+        else {
+            Py_SETREF(obj, (PyObject *) tmp_arr);
+        }
     }
 
     /* Check for Boolean array */
diff --git a/numpy/_core/tests/test_indexing.py b/numpy/_core/tests/test_indexing.py
@@ -622,6 +622,22 @@ def test_nontuple_ndindex(self):
         assert_equal(a[[0, 1], [0, 1]], np.array([0, 6]))
         assert_raises(IndexError, a.__getitem__, [slice(None)])
 
+    def test_flat_index_on_flatiter(self):
+        a = np.arange(9).reshape((3, 3))
+        b = np.array([0, 5, 6])
+        assert_equal(a.flat[b.flat], np.array([0, 5, 6]))
+
+    def test_empty_string_flat_index_on_flatiter(self):
+        a = np.arange(9).reshape((3, 3))
+        b = np.array([], dtype="S")
+        assert_equal(a.flat[b.flat], np.array([]))
+
+    def test_nonempty_string_flat_index_on_flatiter(self):
+        a = np.arange(9).reshape((3, 3))
+        b = np.array(["a"], dtype="S")
+        with pytest.raises(IndexError, match="unsupported iterator index"):
+            a.flat[b.flat]
+
 
 class TestFieldIndexing:
     def test_scalar_return_type(self):
diff --git a/numpy/typing/tests/data/pass/flatiter.py b/numpy/typing/tests/data/pass/flatiter.py
@@ -14,3 +14,6 @@
 a[:]
 a.__array__()
 a.__array__(np.dtype(np.float64))
+
+b = np.array([1]).flat
+a[b]
