BUG: Fix iterator shape in advanced index assignment broadcast error

The reported shape always missed the transpose operation, this refactors
the transpose getting out and applies it in the error path.

Closes gh-18401, gh-5710

2 files changed, 76 insertions, 43 deletions

diff --git a/numpy/core/src/multiarray/mapping.c b/numpy/core/src/multiarray/mapping.c
index 8b9b67387..dd6d6630a 100644
--- a/numpy/core/src/multiarray/mapping.c
+++ b/numpy/core/src/multiarray/mapping.c
@@ -63,15 +63,61 @@ array_length(PyArrayObject *self)
 
 /* -------------------------------------------------------------- */
 
+
+/*
+ * Helper for `PyArray_MapIterSwapAxes` (and related), see its documentation.
+ */
+static void
+_get_transpose(int fancy_ndim, int consec, int ndim, int getmap, npy_intp *dims)
+{
+    /*
+     * For getting the array the tuple for transpose is
+     * (n1,...,n1+n2-1,0,...,n1-1,n1+n2,...,n3-1)
+     * n1 is the number of dimensions of the broadcast index array
+     * n2 is the number of dimensions skipped at the start
+     * n3 is the number of dimensions of the result
+     */
+
+    /*
+     * For setting the array the tuple for transpose is
+     * (n2,...,n1+n2-1,0,...,n2-1,n1+n2,...n3-1)
+     */
+    int n1 = fancy_ndim;
+    int n2 = consec;  /* axes to insert at */
+    int n3 = ndim;
+
+    /* use n1 as the boundary if getting but n2 if setting */
+    int bnd = getmap ? n1 : n2;
+    int val = bnd;
+    int i = 0;
+    while (val < n1 + n2) {
+        dims[i++] = val++;
+    }
+    val = 0;
+    while (val < bnd) {
+        dims[i++] = val++;
+    }
+    val = n1 + n2;
+    while (val < n3) {
+        dims[i++] = val++;
+    }
+}
+
+
 /*NUMPY_API
  *
+ * Swap the axes to or from their inserted form. MapIter always puts the
+ * advanced (array) indices first in the iteration. But if they are
+ * consecutive, will insert/transpose them back before returning.
+ * This is stored as `mit->consec != 0` (the place where they are inserted)
+ * For assignments, the opposite happens: The values to be assigned are
+ * transposed (getmap=1 instead of getmap=0). `getmap=0` and `getmap=1`
+ * undo the other operation.
  */
 NPY_NO_EXPORT void
 PyArray_MapIterSwapAxes(PyArrayMapIterObject *mit, PyArrayObject **ret, int getmap)
 {
     PyObject *new;
-    int n1, n2, n3, val, bnd;
-    int i;
     PyArray_Dims permute;
     npy_intp d[NPY_MAXDIMS];
     PyArrayObject *arr;
@@ -85,10 +131,10 @@ PyArray_MapIterSwapAxes(PyArrayMapIterObject *mit, PyArrayObject **ret, int getm
      */
     arr = *ret;
     if (PyArray_NDIM(arr) != mit->nd) {
-        for (i = 1; i <= PyArray_NDIM(arr); i++) {
+        for (int i = 1; i <= PyArray_NDIM(arr); i++) {
             permute.ptr[mit->nd-i] = PyArray_DIMS(arr)[PyArray_NDIM(arr)-i];
         }
-        for (i = 0; i < mit->nd-PyArray_NDIM(arr); i++) {
+        for (int i = 0; i < mit->nd-PyArray_NDIM(arr); i++) {
             permute.ptr[i] = 1;
         }
         new = PyArray_Newshape(arr, &permute, NPY_ANYORDER);
@@ -99,44 +145,8 @@ PyArray_MapIterSwapAxes(PyArrayMapIterObject *mit, PyArrayObject **ret, int getm
         }
     }
 
-    /*
-     * Setting and getting need to have different permutations.
-     * On the get we are permuting the returned object, but on
-     * setting we are permuting the object-to-be-set.
-     * The set permutation is the inverse of the get permutation.
-     */
-
-    /*
-     * For getting the array the tuple for transpose is
-     * (n1,...,n1+n2-1,0,...,n1-1,n1+n2,...,n3-1)
-     * n1 is the number of dimensions of the broadcast index array
-     * n2 is the number of dimensions skipped at the start
-     * n3 is the number of dimensions of the result
-     */
-
-    /*
-     * For setting the array the tuple for transpose is
-     * (n2,...,n1+n2-1,0,...,n2-1,n1+n2,...n3-1)
-     */
-    n1 = mit->nd_fancy;
-    n2 = mit->consec; /* axes to insert at */
-    n3 = mit->nd;
+    _get_transpose(mit->nd_fancy, mit->consec, mit->nd, getmap, permute.ptr);
 
-    /* use n1 as the boundary if getting but n2 if setting */
-    bnd = getmap ? n1 : n2;
-    val = bnd;
-    i = 0;
-    while (val < n1 + n2) {
-        permute.ptr[i++] = val++;
-    }
-    val = 0;
-    while (val < bnd) {
-        permute.ptr[i++] = val++;
-    }
-    val = n1 + n2;
-    while (val < n3) {
-        permute.ptr[i++] = val++;
-    }
     new = PyArray_Transpose(*ret, &permute);
     Py_DECREF(*ret);
     *ret = (PyArrayObject *)new;
@@ -3200,12 +3210,19 @@ PyArray_MapIterNew(npy_index_info *indices , int index_num, int index_type,
 
     int extra_ndim = PyArray_NDIM(original_extra_op);
     npy_intp *extra_dims = PyArray_DIMS(original_extra_op);
-    PyObject *shape1 = convert_shape_to_string(extra_ndim, extra_dims, " ");
+    PyObject *shape1 = convert_shape_to_string(extra_ndim, extra_dims, "");
     if (shape1 == NULL) {
         goto finish;
     }
 
-    PyObject *shape2 = convert_shape_to_string(mit->nd, mit->dimensions, "");
+    /* Unscramble the iterator shape for reporting when `mit->consec` is used */
+    npy_intp transposed[NPY_MAXDIMS];
+    _get_transpose(mit->nd_fancy, mit->consec, mit->nd, 1, transposed);
+    for (i = 0; i < mit->nd; i++) {
+        transposed[i] = mit->dimensions[transposed[i]];
+    }
+
+    PyObject *shape2 = convert_shape_to_string(mit->nd, transposed, "");
     if (shape2 == NULL) {
         Py_DECREF(shape1);
         goto finish;
diff --git a/numpy/core/tests/test_indexing.py b/numpy/core/tests/test_indexing.py
index 73dbc429c..80f878b35 100644
--- a/numpy/core/tests/test_indexing.py
+++ b/numpy/core/tests/test_indexing.py
@@ -609,6 +609,22 @@ class TestBroadcastedAssignments:
         assert_raises(ValueError, assign, a, s_[:, [0]], np.zeros((5, 0)))
         assert_raises(ValueError, assign, a, s_[[0], :], np.zeros((2, 1)))
 
+    @pytest.mark.parametrize("index", [
+            (..., [1, 2], slice(None)),
+            ([0, 1], ..., 0),
+            (..., [1, 2], [1, 2])])
+    def test_broadcast_error_reports_correct_shape(self, index):
+        values = np.zeros((100, 100))  # will never broadcast below  
+
+        arr = np.zeros((3, 4, 5, 6, 7))
+        # We currently report without any spaces (could be changed)
+        shape_str = str(arr[index].shape).replace(" ", "")
+        
+        with pytest.raises(ValueError) as e:
+            arr[index] = values
+
+        assert str(e.value).endswith(shape_str)
+
     def test_index_is_larger(self):
         # Simple case of fancy index broadcasting of the index.
         a = np.zeros((5, 5))