ENH,TST: Add many tests and have boolean finalize at end.

__array_finalize__ is now called in boolean indexing on the
result array, so that all numerical values are available.

2 files changed, 137 insertions, 44 deletions

diff --git a/numpy/core/src/multiarray/mapping.c b/numpy/core/src/multiarray/mapping.c
index 4dc8bd066..3cd13c1ce 100644
--- a/numpy/core/src/multiarray/mapping.c
+++ b/numpy/core/src/multiarray/mapping.c
@@ -523,7 +523,7 @@ prepare_index(PyArrayObject *self, PyObject *index,
                             "in the future, boolean array-likes will be "
                             "handled as a boolean array index") < 0) {
                         Py_DECREF(tmp_arr);
-                        goto failed_building_indices;  
+                        goto failed_building_indices;
                     }
                     if (PyArray_NDIM(tmp_arr) == 0) {
                         /*
@@ -562,7 +562,7 @@ prepare_index(PyArrayObject *self, PyObject *index,
                                 "numpy.newaxis (`None`) and integer or boolean "
                                 "arrays are valid indices");
                             Py_DECREF((PyObject *)tmp_arr);
-                            goto failed_building_indices; 
+                            goto failed_building_indices;
                         }
                     }
                     else {
@@ -576,9 +576,9 @@ prepare_index(PyArrayObject *self, PyObject *index,
                                 "non integer (and non boolean) array-likes will "
                                 "not be accepted as indices in the future");
                             Py_DECREF((PyObject *)tmp_arr);
-                            goto failed_building_indices;  
+                            goto failed_building_indices;
                         }
-                    }   
+                    }
                 }
 
                 PyArray_Descr *indtype = PyArray_DescrFromType(NPY_INTP);
@@ -984,8 +984,8 @@ array_boolean_subscript(PyArrayObject *self,
     /* Allocate the output of the boolean indexing */
     dtype = PyArray_DESCR(self);
     Py_INCREF(dtype);
-    ret = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(self), dtype, 1, &size,
-                                NULL, NULL, 0, (PyObject *)self);
+    ret = (PyArrayObject *)PyArray_NewFromDescr(&PyArray_Type, dtype, 1, &size,
+                                NULL, NULL, 0, NULL);
     if (ret == NULL) {
         return NULL;
     }
@@ -1074,6 +1074,24 @@ array_boolean_subscript(PyArrayObject *self,
         NPY_AUXDATA_FREE(transferdata);
     }
 
+    if (!PyArray_CheckExact(self)) {
+        PyArrayObject *tmp = ret;
+
+        ret = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(self), dtype, 1,
+                            &size, PyArray_STRIDES(ret), PyArray_BYTES(ret),
+                            0, (PyObject *)self);
+
+        if (ret == NULL) {
+            Py_DECREF(tmp);
+            return NULL;
+        }
+
+        if (PyArray_SetBaseObject(ret, tmp) < 0) {
+            Py_DECREF(ret);
+            return NULL;
+        }
+    }
+
     return ret;
 }
 
@@ -1511,10 +1529,11 @@ array_subscript(PyArrayObject *self, PyObject *op)
             Py_DECREF(tmp_arr);
             goto finish;
         }
-        /* TODO: Could attempt to steal the data from the old result */
+
         if (PyArray_SetBaseObject((PyArrayObject *)result,
                                   (PyObject *)tmp_arr) < 0) {
-            Py_DECREF(tmp_arr);
+            Py_DECREF(result);
+            result = NULL;
             goto finish;
         }
     }
@@ -1686,15 +1705,6 @@ array_ass_sub(PyArrayObject *self, PyObject *ind, PyObject *op)
 
     /* If there is no fancy indexing, we have the array to assign to */
     if (!(index_type & HAS_FANCY)) {
-        /*
-         * CopyObject handles all weirdness for us, this however also
-         * means that other array assignments which convert more strictly
-         * do *not* handle all weirdnesses correctly.
-         * TODO: To have other assignments handle them correctly, we
-         *       should copy into a temporary array of the correct shape
-         *       if it is not an array yet!
-         * TODO: We could use PyArray_SETITEM if it is 0-d?
-         */
         if (PyArray_CopyObject(view, op) < 0) {
             goto fail;
         }
@@ -1702,9 +1712,13 @@ array_ass_sub(PyArrayObject *self, PyObject *ind, PyObject *op)
     }
 
     if (!PyArray_Check(op)) {
-        if ((PyDataType_HASFIELDS(descr) || PyDataType_REFCHK(descr))
-                    && PySequence_Check(op)) {
-            /* Allocate array through MapIter to fill with CopyObject */
+        /*
+         * If the array is of object converting the values to an array
+         * might not be legal even though normal assignment works.
+         * So allocate a temporary array of the right size and use the
+         * normal assignment to handle this case.
+         */
+        if (PyDataType_REFCHK(descr) && PySequence_Check(op)) {
             tmp_arr = NULL;
         }
         else {
@@ -2446,17 +2460,16 @@ PyArray_MapIterNew(npy_index_info *indices , int index_num, int index_type,
      *   2. Subspace iteration is necessary, so the extra op is iterated
      *      independendly, and the iteration order is fixed at C (could
      *      also use Fortran order if the array is Fortran order).
-     *      In this case the subspace iterator is buffered, so that even
-     *      the array being iterated *is* buffered!
+     *      In this case the subspace iterator is not buffered.
      *   3. Subspace iteration is necessary and an extra_op was given.
-     *      In this case it needs to be transposed!
+     *      In this case it may need transposing, etc.
      */
 
     /*
      * If we have an extra_op given, need to prepare it.
      *   1. Subclasses might mess with the shape, so need a baseclass
      *   2. Need to make sure the shape is compatible
-     *   3. May need to transpose it.
+     *   3. May need to remove leading 1s and transpose dimensions.
      */
     if (extra_op != NULL) {
         if (!PyArray_CheckExact(extra_op)) {
@@ -2471,10 +2484,9 @@ PyArray_MapIterNew(npy_index_info *indices , int index_num, int index_type,
 
         if (PyArray_NDIM(extra_op) > mit->nd) {
             /*
-             * Usual assignments allows removal of trailing one dimensions.
+             * Usual assignments allows removal of leading 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.
+             * assigned to). To implement this, reshape the array.
              * This should maybe be done differently, or even not be allowed.
              */
             PyArrayObject *tmp_arr;
diff --git a/numpy/core/tests/test_indexing.py b/numpy/core/tests/test_indexing.py
index ab68326c8..892a85811 100644
--- a/numpy/core/tests/test_indexing.py
+++ b/numpy/core/tests/test_indexing.py
@@ -147,11 +147,59 @@ class TestIndexing(TestCase):
                          [0, 8, 0]])
 
 
+    def test_reverse_strides_and_subspace_bufferinit(self):
+        # 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(self):
+        # Just documents behaviour, this is a small limitation.
+        a = np.ones((1,) * 32) # 32 is NPY_MAXDIMS
+        assert_raises(IndexError, a.__getitem__, (np.array([0]),) * 32)
+
+
+    def test_scalar_array_bool(self):
+        # Numpy bools can be used as boolean index (python ones as of yet not)
+        a = np.array(1)
+        assert_equal(a[np.bool_(True)], a[np.array(True)])
+        assert_equal(a[np.bool_(False)], a[np.array(False)])
+
+        # After deprecating bools as integers:
+        #a = np.array([0,1,2])
+        #assert_equal(a[True, :], a[None, :])
+        #assert_equal(a[:, True], a[:, None])
+        #
+        #assert_(not np.may_share_memory(a, a[True, :]))
+
+
+    def test_everything_returns_views(self):
+        # Before `...` would return a itself.
+        a = np.arange(5)
+
+        assert_(a is not a[()])
+        assert_(a is not a[...])
+        assert_(a is not a[:])
+
+
 class TestBroadcastedAssignments(TestCase):
     def assign(self, a, ind, val):
         a[ind] = val
         return a
 
+
     def test_prepending_ones(self):
         a = np.zeros((3, 2))
 
@@ -189,27 +237,60 @@ class TestBroadcastedAssignments(TestCase):
         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]
+    def test_index_is_larger(self):
+        # Simple case of fancy index broadcastin of the index.
+        a = np.zeros((5, 5))
+        a[[[0], [1], [2]], [0, 1, 2]] = [2, 3, 4]
+
+        assert_((a[:3, :3] == [2, 3, 4]).all())
+
+
+class TestSubclasses(TestCase):
+    def test_basic(self):
+        class SubClass(np.ndarray):
+            pass
+
+        s = np.arange(5).view(SubClass)
+        assert_(isinstance(s[:3], SubClass))
+        assert_(s[:3].base is s)
+
+        assert_(isinstance(s[[0, 1, 2]], SubClass))
+        assert_(isinstance(s[s > 0], SubClass))
+
+
+    def test_matrix_fancy(self):
+        # The matrix class messes with the shape. While this is always
+        # weird (getitem is not used, it does not have setitem nor knows
+        # about fancy indexing), this tests gh-3110
+        m = np.matrix([[1, 2], [3, 4]])
+
+        assert_(isinstance(m[[0,1,0], :], np.matrix))
+
+        # gh-3110. Note the transpose currently because matrices do *not*
+        # support dimension fixing for fancy indexing correctly.
+        x = np.asmatrix(np.arange(50).reshape(5,10))
+        assert_equal(x[:2, np.array(-1)], x[:2, -1].T)
+
 
-    a[b] = c
-    # If the strides are not reversed, the 0 in the arange comes last.
-    assert_equal(a[0], 0)
+    def test_finalize_gets_full_info(self):
+        # Array finalize should be called on the filled array.
+        class SubClass(np.ndarray):
+            def __array_finalize__(self, old):
+                self.finalize_status = np.array(self)
+                self.old = old
 
-    # 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])
+        s = np.arange(10).view(SubClass)
+        new_s = s[:3]
+        assert_array_equal(new_s.finalize_status, new_s)
+        assert_array_equal(new_s.old, s)
 
+        new_s = s[[0,1,2,3]]
+        assert_array_equal(new_s.finalize_status, new_s)
+        assert_array_equal(new_s.old, s)
 
-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)
+        new_s = s[s > 0]
+        assert_array_equal(new_s.finalize_status, new_s)
+        assert_array_equal(new_s.old, s)
 
 
 class TestFancyIndexingEquivalence(TestCase):