BUGS: Fix subspace casting and make fancy broadcasting consistent

Normal CopyInto allows removal of 1s from the start, but
broadcasting never did (and did not here before). This allows that
case, though by creating a view.

3 files changed, 139 insertions, 63 deletions

diff --git a/numpy/core/src/multiarray/lowlevel_strided_loops.c.src b/numpy/core/src/multiarray/lowlevel_strided_loops.c.src
index de64965ef..fcdb2d8ed 100644
--- a/numpy/core/src/multiarray/lowlevel_strided_loops.c.src
+++ b/numpy/core/src/multiarray/lowlevel_strided_loops.c.src
@@ -1614,13 +1614,15 @@ mapiter_@name@(PyArrayMapIterObject *mit)
          * buffers, this is safe.
          */
         NpyIter_GetInnerFixedStrideArray(mit->subspace_iter, fixed_strides);
+
         if (PyArray_GetDTypeTransferFunction(is_aligned,
 #if @isget@
                         fixed_strides[0], fixed_strides[1],
+                        PyArray_DESCR(array), PyArray_DESCR(mit->extra_op),
 #else
                         fixed_strides[1], fixed_strides[0],
+                         PyArray_DESCR(mit->extra_op), PyArray_DESCR(array),
 #endif
-                        PyArray_DESCR(array), PyArray_DESCR(mit->extra_op),
                         0,
                         &stransfer, &transferdata,
                         &needs_api) != NPY_SUCCEED) {
@@ -1710,10 +1712,10 @@ mapiter_@name@(PyArrayMapIterObject *mit)
 
                 mit->extra_op_next(mit->extra_op_iter);
             } while (mit->outer_next(mit->outer));
-
-            NPY_AUXDATA_FREE(transferdata);
         }
 /**end repeat1**/
+
+        NPY_AUXDATA_FREE(transferdata);
     }
     return 0;
 }
diff --git a/numpy/core/src/multiarray/mapping.c b/numpy/core/src/multiarray/mapping.c
index ee8d8eef4..4dc8bd066 100644
--- a/numpy/core/src/multiarray/mapping.c
+++ b/numpy/core/src/multiarray/mapping.c
@@ -1817,6 +1817,7 @@ array_ass_sub(PyArrayObject *self, PyObject *ind, PyObject *op)
         Py_XDECREF(indices[i].object);
     }
     return -1;
+
   success:
     Py_XDECREF((PyObject *)view);
     Py_XDECREF((PyObject *)tmp_arr);
@@ -2355,6 +2356,8 @@ PyArray_MapIterNew(npy_index_info *indices , int index_num, int index_type,
                    PyArray_Descr *extra_op_dtype)
 {
     PyObject *errmsg, *tmp;
+    /* For shape reporting on error */
+    PyArrayObject *original_extra_op = extra_op;
 
     PyArrayObject *index_arrays[NPY_MAXDIMS];
     PyArray_Descr *dtypes[NPY_MAXDIMS];
@@ -2467,8 +2470,25 @@ PyArray_MapIterNew(npy_index_info *indices , int index_num, int index_type,
         }
 
         if (PyArray_NDIM(extra_op) > mit->nd) {
-            /* TODO: Add a test. */
-            goto broadcast_error;
+            /*
+             * Usual assignments allows removal of trailing one dimensions.
+             * (or equivalently adding of one dimensions to the array being
+             * assigned to). To implement this, reshape the array. It is a bit
+             * of a hack.
+             * This should maybe be done differently, or even not be allowed.
+             */
+            PyArrayObject *tmp_arr;
+            PyArray_Dims permute;
+
+            permute.len = mit->nd;
+            permute.ptr = &PyArray_DIMS(extra_op)[
+                                            PyArray_NDIM(extra_op) - mit->nd];
+            tmp_arr = PyArray_Newshape(extra_op, &permute, NPY_CORDER);
+            if (tmp_arr == NULL) {
+                goto broadcast_error;
+            }
+            Py_DECREF(extra_op);
+            extra_op = tmp_arr;
         }
 
         /*
@@ -2630,7 +2650,7 @@ PyArray_MapIterNew(npy_index_info *indices , int index_num, int index_type,
                          "when no subspace is given, the number of index "
                          "arrays cannot be above %d, but %d index arrays found",
                          NPY_MAXDIMS - 1, mit->numiter);
-            return NULL;
+            goto fail;
         }
 
         nops += 1;
@@ -2695,19 +2715,17 @@ PyArray_MapIterNew(npy_index_info *indices , int index_num, int index_type,
     if (extra_op_flags) {
         if (extra_op == NULL) {
             mit->extra_op = NpyIter_GetOperandArray(mit->outer)[mit->numiter];
-            Py_INCREF(mit->extra_op);
         }
         else {
             mit->extra_op = extra_op;
         }
+        Py_INCREF(mit->extra_op);
     }
 
     /*
      * If extra_op is being tracked but subspace is used, we need
      * to create a dedicated iterator for the outer iteration of
      * the extra operand.
-     *
-     * Also create the external op offset iteration.
      */
     if (extra_op_flags && uses_subspace) {
         op_axes[0] = single_op_axis;
@@ -2751,6 +2769,7 @@ PyArray_MapIterNew(npy_index_info *indices , int index_num, int index_type,
 
     /* Can now return early if no subspace is being used */
     if (!uses_subspace) {
+        Py_XDECREF(extra_op);
         return (PyObject *)mit;
     }
 
@@ -2789,17 +2808,17 @@ PyArray_MapIterNew(npy_index_info *indices , int index_num, int index_type,
     }
 
     mit->subspace_iter = NpyIter_AdvancedNew(nops, index_arrays,
-                                        NPY_ITER_ZEROSIZE_OK |
-                                        NPY_ITER_REFS_OK |
-                                        NPY_ITER_GROWINNER |
-                                        NPY_ITER_EXTERNAL_LOOP |
-                                        NPY_ITER_DELAY_BUFALLOC |
-                                        subspace_iter_flags,
-                                        (nops == 1 ? NPY_CORDER : NPY_KEEPORDER),
-                                        NPY_UNSAFE_CASTING,
-                                        op_flags, dtypes,
-                                        PyArray_NDIM(subspace), op_axes,
-                                        &mit->dimensions[mit->nd_fancy], 0);
+                                    NPY_ITER_ZEROSIZE_OK |
+                                    NPY_ITER_REFS_OK |
+                                    NPY_ITER_GROWINNER |
+                                    NPY_ITER_EXTERNAL_LOOP |
+                                    NPY_ITER_DELAY_BUFALLOC |
+                                    subspace_iter_flags,
+                                    (nops == 1 ? NPY_CORDER : NPY_KEEPORDER),
+                                    NPY_UNSAFE_CASTING,
+                                    op_flags, dtypes,
+                                    PyArray_NDIM(subspace), op_axes,
+                                    &mit->dimensions[mit->nd_fancy], 0);
 
     if (mit->subspace_iter == NULL) {
         goto fail;
@@ -2820,72 +2839,74 @@ PyArray_MapIterNew(npy_index_info *indices , int index_num, int index_type,
          */
     }
 
+    Py_XDECREF(extra_op);
     return (PyObject *)mit;
 
   fail:
     /*
      * Check whether the operand was not broadcastable and replace the error
-     * in that case, since the NpyIter error speaks of iterators and
-     * remapped axes, etc...
+     * in that case. This should however normally be found early with a
+     * direct goto to broadcast_error
      */
     if (extra_op == NULL) {
-        goto no_broadcast_error;
+        goto finish;
     }
 
     j = mit->nd;
     for (i = PyArray_NDIM(extra_op) - 1; i >= 0; i--) {
         j--;
-        if ((PyArray_DIM(extra_op, i) != 1)
-                &&  (PyArray_DIM(extra_op, i) != mit->dimensions[j])) {
+        if ((PyArray_DIM(extra_op, i) != 1) &&
+                /* (j < 0 is currently impossible, extra_op is reshaped) */
+                j >= 0 &&
+                PyArray_DIM(extra_op, i) != mit->dimensions[j]) {
             /* extra_op cannot be broadcasted to the indexing result */
             goto broadcast_error;
         }
     }
-    goto no_broadcast_error;
+    goto finish;
 
   broadcast_error:
     errmsg = PyUString_FromString("shape mismatch: value array "
                     "of shape ");
     if (errmsg == NULL) {
-        Py_DECREF(mit);
-        return NULL;
+        goto finish;
     }
 
-    tmp = get_shape_string(PyArray_NDIM(extra_op),
-                           PyArray_DIMS(extra_op), " ");
+    /* Report the shape of the original array if it exists */
+    if (original_extra_op == NULL) {
+        original_extra_op = extra_op;
+    }
+
+    tmp = get_shape_string(PyArray_NDIM(original_extra_op),
+                           PyArray_DIMS(original_extra_op), " ");
     if (tmp == NULL) {
-        Py_DECREF(mit);
-        return NULL;
+        goto finish;
     }
     PyUString_ConcatAndDel(&errmsg, tmp);
     if (errmsg == NULL) {
-        Py_DECREF(mit);
-        return NULL;
+        goto finish;
     }
 
     tmp = PyUString_FromString("could not be broadcast to indexing "
                     "result of shape ");
     PyUString_ConcatAndDel(&errmsg, tmp);
     if (errmsg == NULL) {
-        Py_DECREF(mit);
-        return NULL;
+        goto finish;
     }
 
     tmp = get_shape_string(mit->nd, mit->dimensions, "");
     if (tmp == NULL) {
-        Py_DECREF(mit);
-        return NULL;
+        goto finish;
     }
     PyUString_ConcatAndDel(&errmsg, tmp);
     if (errmsg == NULL) {
-        Py_DECREF(mit);
-        return NULL;
+        goto finish;
     }
 
     PyErr_SetObject(PyExc_ValueError, errmsg);
     Py_DECREF(errmsg);
 
-  no_broadcast_error:
+  finish:
     Py_XDECREF(extra_op);
     Py_DECREF(mit);
     return NULL;
diff --git a/numpy/core/tests/test_indexing.py b/numpy/core/tests/test_indexing.py
index b1dd2195b..ab68326c8 100644
--- a/numpy/core/tests/test_indexing.py
+++ b/numpy/core/tests/test_indexing.py
@@ -150,44 +150,97 @@ class TestIndexing(TestCase):
 class TestBroadcastedAssignments(TestCase):
     def assign(self, a, ind, val):
         a[ind] = val
+        return a
 
+    def test_prepending_ones(self):
+        a = np.zeros((3, 2))
 
-    def test_non_fancy(self):
-        assign = self.assign
-        s_ = np.s_
+        a[...] = np.ones((1, 3, 2))
+        # Fancy with subspace with and without transpose
+        a[[0, 1, 2], :] = np.ones((1, 3, 2))
+        a[:, [0, 1]] = np.ones((1, 3, 2))
+        # Fancy without subspace (with broadcasting)
+        a[[[0], [1], [2]], [0, 1]] = np.ones((1, 3, 2))
 
-        a = np.zeros((1,2,0))
 
-        assert_raises(ValueError, assign, a, s_[...], np.ones((1,2,3)))
-        # Too large should raise
-        assert_raises(ValueError, assign, a, s_[...], np.ones((1,1,1,1)))
-        a[...] = np.ones((1,1,1))
-
-    def test_fancy_no_subspace(self):
+    def test_prepend_not_one(self):
         assign = self.assign
         s_ = np.s_
 
         a = np.zeros(5)
 
-        a[[1,2,3]] = 0
-        assert_raises(ValueError, a, s_[[1,2,3]], [1,2,3,4])
+        # Too large and not only ones.
+        assert_raises(ValueError, assign, a, s_[...],  np.ones((2, 1)))
+        assert_raises(ValueError, assign, a, s_[[1, 2, 3],],  np.ones((2, 1)))
+        assert_raises(ValueError, assign, a, s_[[[1], [2]],], np.ones((2,2,1)))
 
-        # values can be extended, result cannot
-        a[[[1,2,3]]] = [1,2,3]
-        assert_raises(ValueError, a, s_[[1,2,3]], [[1,2,3]])
 
-    def test_fancy_with_ignored_subspace(self):
+    def test_simple_broadcasting_errors(self):
         assign = self.assign
         s_ = np.s_
 
-        a = np.zeros(1,5,1)
+        a = np.zeros((5, 1))
+        assert_raises(ValueError, assign, a, s_[...], np.zeros((5, 2)))
+        assert_raises(ValueError, assign, a, s_[...], np.zeros((5, 0)))
+
+        assert_raises(ValueError, assign, a, s_[:, [0]], np.zeros((5, 2)))
+        assert_raises(ValueError, assign, a, s_[:, [0]], np.zeros((5, 0)))
+
+        assert_raises(ValueError, assign, a, s_[[0], :], np.zeros((2, 1)))
+
+
+def test_reverse_strides_and_subspace_bufferinit():
+    # This tests that the strides are not reversed for simple and
+    # subspace fancy indexing.
+    a = np.ones(5)
+    b = np.zeros(5, dtype=np.intp)[::-1]
+    c = np.arange(5)[::-1]
+
+    a[b] = c
+    # If the strides are not reversed, the 0 in the arange comes last.
+    assert_equal(a[0], 0)
+
+    # This also tests that the subspace buffer is initiliazed:
+    a = np.ones((5, 2))
+    c = np.arange(10).reshape(5, 2)[::-1]
+    a[b, :] = c
+    assert_equal(a[0], [0, 1])
+
+
+def test_too_many_fancy_indices_special_case():
+    a = np.ones((1,) * 32) # 32 is NPY_MAXDIMS
+    assert_raises(IndexError, a.__getitem__, (np.array([0]),) * 32)
+
+
+class TestFancyIndexingEquivalence(TestCase):
+    def test_object_assign(self):
+        # Check that the field and object special case using copyto is active
+        # The right hand side cannot be converted to an array here.
+        a = np.arange(5, dtype=object)
+        b = a.copy()
+        a[:3] = [1, (1,2), 3]
+        b[[0, 1, 2]] = [1, (1,2), 3]
+        assert_array_equal(a, b)
+
+        # test same for subspace fancy indexing
+        b = np.arange(5, dtype=object)[None, :]
+        b[[0], :3] = [[1, (1,2), 3]]
+        assert_array_equal(a, b[0])
+
+
+    def test_cast_equivalence(self):
+        # Yes, normal slicing uses unsafe casting.
+        a = np.arange(5)
+        b = a.copy()
 
-        a[:,[1,2,3],:] = [1,2,3]
-        assert_array_equal(a, [[[1], [2], [3]]])
-        assert_raises(ValueError, assign, a, s_[:,[1],:], np.ones((1,) * 4))
+        a[:3] = np.array(['2', '-3', '-1'])
+        b[[0, 2, 1]] = np.array(['2', '-1', '-3'])
+        assert_array_equal(a, b)
 
-    def test_fancy_with_subspace(self):
-        pass
+        # test the same for subspace fancy indexing
+        b = np.arange(5)[None, :]
+        b[[0], :3] = np.array([['2', '-3', '-1']])
+        assert_array_equal(a, b[0])
 
 
 class TestMultiIndexingAutomated(TestCase):