fix from review

2 files changed, 17 insertions, 15 deletions

diff --git a/numpy/core/include/numpy/ndarrayobject.h b/numpy/core/include/numpy/ndarrayobject.h
index 97e41b6f3..f17008b0b 100644
--- a/numpy/core/include/numpy/ndarrayobject.h
+++ b/numpy/core/include/numpy/ndarrayobject.h
@@ -174,15 +174,12 @@ extern "C" CONFUSE_EMACS
 static NPY_INLINE void
 PyArray_DiscardWritebackIfCopy(PyArrayObject *arr)
 {
-    if (arr != NULL) {
-        PyArrayObject_fields *fa = (PyArrayObject_fields *)arr;
-        if ((PyArray_FLAGS(arr) & NPY_ARRAY_WRITEBACKIFCOPY) ||
-            (PyArray_FLAGS(arr) & NPY_ARRAY_UPDATEIFCOPY)) {
-            if (fa->base) {
-                PyArray_ENABLEFLAGS((PyArrayObject*)fa->base, NPY_ARRAY_WRITEABLE);
-                Py_DECREF(fa->base);
-                fa->base = NULL;
-            }
+    PyArrayObject_fields *fa = (PyArrayObject_fields *)arr;
+    if (fa && fa->base) {
+        if ((fa->flags & NPY_ARRAY_UPDATEIFCOPY) || (fa->flags & NPY_ARRAY_WRITEBACKIFCOPY)) {
+            PyArray_ENABLEFLAGS((PyArrayObject*)fa->base, NPY_ARRAY_WRITEABLE);
+            Py_DECREF(fa->base);
+            fa->base = NULL;
             PyArray_CLEARFLAGS(arr, NPY_ARRAY_WRITEBACKIFCOPY);
             PyArray_CLEARFLAGS(arr, NPY_ARRAY_UPDATEIFCOPY);
         }
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index 16d47839f..ed3102f66 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -7246,17 +7246,20 @@ class TestWritebackIfCopy(object):
 
     def test_view_assign(self):
         from numpy.core._multiarray_tests import npy_create_writebackifcopy, npy_resolve
+
         arr = np.arange(9).reshape(3, 3).T
         arr_wb = npy_create_writebackifcopy(arr)
         assert_(arr_wb.flags.writebackifcopy)
         assert_(arr_wb.base is arr)
-        arr_wb[:] = -100
+        arr_wb[...] = -100
         npy_resolve(arr_wb)
+        # arr changes after resolve, even though we assigned to arr_wb
         assert_equal(arr, -100)
         # after resolve, the two arrays no longer reference each other
-        assert_(not arr_wb.ctypes.data == 0)
+        assert_(arr_wb.ctypes.data != 0)
         assert_equal(arr_wb.base, None)
-        arr_wb[:] = 100
+        # assigning to arr_wb does not get transfered to arr
+        arr_wb[...] = 100
         assert_equal(arr, -100)
 
     def test_dealloc_warning(self):
@@ -7269,6 +7272,7 @@ class TestWritebackIfCopy(object):
 
     def test_view_discard_refcount(self):
         from numpy.core._multiarray_tests import npy_create_writebackifcopy, npy_discard
+
         arr = np.arange(9).reshape(3, 3).T
         orig = arr.copy()
         if HAS_REFCOUNT:
@@ -7276,16 +7280,17 @@ class TestWritebackIfCopy(object):
         arr_wb = npy_create_writebackifcopy(arr)
         assert_(arr_wb.flags.writebackifcopy)
         assert_(arr_wb.base is arr)
-        arr_wb[:] = -100
+        arr_wb[...] = -100
         npy_discard(arr_wb)
         # arr remains unchanged after discard
         assert_equal(arr, orig)
         # after discard, the two arrays no longer reference each other
-        assert_(not arr_wb.ctypes.data == 0)
+        assert_(arr_wb.ctypes.data != 0)
         assert_equal(arr_wb.base, None)
         if HAS_REFCOUNT:
             assert_equal(arr_cnt, sys.getrefcount(arr))
-        arr_wb[:] = 100
+        # assigning to arr_wb does not get transfered to arr
+        arr_wb[...] = 100
         assert_equal(arr, orig)