Merge pull request #10824 from mattip/fix-PyArray_DeprecateWritebackIfCopy

BUG: test, fix PyArray_DiscardWritebackIfCopy refcount issue and docu…

4 files changed, 69 insertions, 19 deletions

diff --git a/doc/source/reference/c-api.array.rst b/doc/source/reference/c-api.array.rst
index ad7c725a8..5ea7bfcfc 100644
--- a/doc/source/reference/c-api.array.rst
+++ b/doc/source/reference/c-api.array.rst
@@ -1360,7 +1360,7 @@ Special functions for NPY_OBJECT
 .. c:function:: int PyArray_SetWritebackIfCopyBase(PyArrayObject* arr, PyArrayObject* base)
 
     Precondition: ``arr`` is a copy of ``base`` (though possibly with different
-    strides, ordering, etc.) Sets the :c:data:`NPY_ARRAY_WRITEBACKIFCOPY` flag 
+    strides, ordering, etc.) Sets the :c:data:`NPY_ARRAY_WRITEBACKIFCOPY` flag
     and ``arr->base``, and set ``base`` to READONLY. Call
     :c:func:`PyArray_ResolveWritebackIfCopy` before calling
     `Py_DECREF`` in order copy any changes back to ``base`` and
@@ -3260,12 +3260,14 @@ Memory management
 .. c:function:: int PyArray_ResolveWritebackIfCopy(PyArrayObject* obj)
 
     If ``obj.flags`` has :c:data:`NPY_ARRAY_WRITEBACKIFCOPY` or (deprecated)
-    :c:data:`NPY_ARRAY_UPDATEIFCOPY`, this function copies ``obj->data`` to
-    `obj->base->data`, clears the flags, `DECREF` s `obj->base` and makes it
-    writeable, and sets ``obj->base`` to NULL. This is the opposite of 
+    :c:data:`NPY_ARRAY_UPDATEIFCOPY`, this function clears the flags, `DECREF` s
+    `obj->base` and makes it writeable, and sets ``obj->base`` to NULL. It then
+    copies ``obj->data`` to `obj->base->data`, and returns the error state of
+    the copy operation. This is the opposite of
     :c:func:`PyArray_SetWritebackIfCopyBase`. Usually this is called once
     you are finished with ``obj``, just before ``Py_DECREF(obj)``. It may be called
-    multiple times, or with ``NULL`` input.
+    multiple times, or with ``NULL`` input. See also
+    :c:func:`PyArray_DiscardWritebackIfCopy`.
 
     Returns 0 if nothing was done, -1 on error, and 1 if action was taken.
 
@@ -3487,12 +3489,14 @@ Miscellaneous Macros
 
 .. c:function:: PyArray_DiscardWritebackIfCopy(PyObject* obj)
 
-    Reset the :c:data:`NPY_ARRAY_WRITEBACKIFCOPY` and deprecated
-    :c:data:`NPY_ARRAY_UPDATEIFCOPY` flag. Resets the
-    :c:data:`NPY_ARRAY_WRITEABLE` flag on the base object. It also
-    discards pending changes to the base object. This is
-    useful for recovering from an error condition when
-    writeback semantics are used.
+    If ``obj.flags`` has :c:data:`NPY_ARRAY_WRITEBACKIFCOPY` or (deprecated)
+    :c:data:`NPY_ARRAY_UPDATEIFCOPY`, this function clears the flags, `DECREF` s
+    `obj->base` and makes it writeable, and sets ``obj->base`` to NULL. In
+    contrast to :c:func:`PyArray_DiscardWritebackIfCopy` it makes no attempt
+    to copy the data from `obj->base` This undoes
+    :c:func:`PyArray_SetWritebackIfCopyBase`. Usually this is called after an
+    error when you are finished with ``obj``, just before ``Py_DECREF(obj)``.
+    It may be called multiple times, or with ``NULL`` input.
 
 .. c:function:: PyArray_XDECREF_ERR(PyObject* obj)
 
diff --git a/numpy/core/include/numpy/ndarrayobject.h b/numpy/core/include/numpy/ndarrayobject.h
index ec0fd1ee9..12fc7098c 100644
--- a/numpy/core/include/numpy/ndarrayobject.h
+++ b/numpy/core/include/numpy/ndarrayobject.h
@@ -170,14 +170,17 @@ extern "C" CONFUSE_EMACS
                                             (k)*PyArray_STRIDES(obj)[2] + \
                                             (l)*PyArray_STRIDES(obj)[3]))
 
+/* Move to arrayobject.c once PyArray_XDECREF_ERR is removed */
 static NPY_INLINE void
 PyArray_DiscardWritebackIfCopy(PyArrayObject *arr)
 {
-    if (arr != NULL) {
-        if ((PyArray_FLAGS(arr) & NPY_ARRAY_WRITEBACKIFCOPY) ||
-            (PyArray_FLAGS(arr) & NPY_ARRAY_UPDATEIFCOPY)) {
-            PyArrayObject *base = (PyArrayObject *)PyArray_BASE(arr);
-            PyArray_ENABLEFLAGS(base, NPY_ARRAY_WRITEABLE);
+    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/src/multiarray/_multiarray_tests.c.src b/numpy/core/src/multiarray/_multiarray_tests.c.src
index 38698887a..0299f1a1b 100644
--- a/numpy/core/src/multiarray/_multiarray_tests.c.src
+++ b/numpy/core/src/multiarray/_multiarray_tests.c.src
@@ -687,6 +687,18 @@ npy_resolve(PyObject* NPY_UNUSED(self), PyObject* args)
     Py_RETURN_NONE;
 }
 
+/* resolve WRITEBACKIFCOPY */
+static PyObject*
+npy_discard(PyObject* NPY_UNUSED(self), PyObject* args)
+{
+    if (!PyArray_Check(args)) {
+        PyErr_SetString(PyExc_TypeError, "test needs ndarray input");
+        return NULL;
+    }
+    PyArray_DiscardWritebackIfCopy((PyArrayObject*)args);
+    Py_RETURN_NONE;
+}
+
 #if !defined(NPY_PY3K)
 static PyObject *
 int_subclass(PyObject *dummy, PyObject *args)
@@ -1857,6 +1869,9 @@ static PyMethodDef Multiarray_TestsMethods[] = {
     {"npy_resolve",
         npy_resolve,
         METH_O, NULL},
+    {"npy_discard",
+        npy_discard,
+        METH_O, NULL},
 #if !defined(NPY_PY3K)
     {"test_int_subclass",
         int_subclass,
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index 806a3b083..ed3102f66 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -7246,16 +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)
-        arr_wb[:] = 100
+        assert_(arr_wb.ctypes.data != 0)
+        assert_equal(arr_wb.base, None)
+        # assigning to arr_wb does not get transfered to arr
+        arr_wb[...] = 100
         assert_equal(arr, -100)
 
     def test_dealloc_warning(self):
@@ -7266,6 +7270,30 @@ class TestWritebackIfCopy(object):
             _multiarray_tests.npy_abuse_writebackifcopy(v)
             assert len(sup.log) == 1
 
+    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:
+            arr_cnt = sys.getrefcount(arr)
+        arr_wb = npy_create_writebackifcopy(arr)
+        assert_(arr_wb.flags.writebackifcopy)
+        assert_(arr_wb.base is arr)
+        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_(arr_wb.ctypes.data != 0)
+        assert_equal(arr_wb.base, None)
+        if HAS_REFCOUNT:
+            assert_equal(arr_cnt, sys.getrefcount(arr))
+        # assigning to arr_wb does not get transfered to arr
+        arr_wb[...] = 100
+        assert_equal(arr, orig)
+
+
 class TestArange(object):
     def test_infinite(self):
         assert_raises_regex(