numpy · mattip · Jun 18, 2017 · Jun 19, 2017 · Jun 19, 2017 · Jun 19, 2017
diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
@@ -381,6 +381,21 @@ def luf(lamdaexpr, *args, **kwargs):
         >>> luf(lambda i,j:i*i + j/2, a, b)
         array([  0.5,   1.5,   4.5,   9.5,  16.5])
 
+    If operand flags such as "updateifcopy" or "readwrite" are used, or in
+    other edge cases, the operands may be views into the original data with the
+    UPDATEIFCOPY flag. In this case it is important to use the nditer as a
+    context manager. The temporary data will be written back to the original
+    data when the context manager's __exit__ function is called::
+
+        >>> a = np.arange(6, dtype='i4')[::-2]
+        >>> with nditer(a, [],
+        ...        [['writeonly', 'updateifcopy']],
+        ...        casting='unsafe',
+        ...        op_dtypes=[np.dtype('f4')]) as i:
+        ...    i.operands[0][:] = [-1, -2, -3]
+        >>> a
+        array([-1, -2, -3])
+
     """)
 
 # nditer methods

diff --git a/numpy/core/code_generators/cversions.txt b/numpy/core/code_generators/cversions.txt
@@ -36,5 +36,5 @@
 0x0000000a = 9b8bce614655d3eb02acddcb508203cb
 
 # Version 11 (NumPy 1.13) Added PyArray_MapIterArrayCopyIfOverlap
-# Version 11 (NumPy 1.14) No Change
+# Version 11 (NumPy 1.14) Added PyArray_ResolveUpdateIfCopy
 0x0000000b = edb1ba83730c650fd9bc5772a919cda7
diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
@@ -346,6 +346,8 @@
     # End 1.10 API
     'PyArray_MapIterArrayCopyIfOverlap':    (301,),
     # End 1.13 API
+    'PyArray_ResolveUpdateIfCopy':    (302,),
+    # End 1.14 API
 }
 
 ufunc_types_api = {

diff --git a/numpy/core/include/numpy/ndarraytypes.h b/numpy/core/include/numpy/ndarraytypes.h
@@ -1027,6 +1027,8 @@ typedef void (NpyIter_GetMultiIndexFunc)(NpyIter *iter,
  * eliminate overlap.
  */
 #define NPY_ITER_COPY_IF_OVERLAP            0x00002000
+/* At least one operand uses UPDATEIFCOPY semantics, use context manager */
+#define NPY_ITER_USE_CONTEXT_MANAGER        0x00004000
 
 /*** Per-operand flags that may be passed to the iterator constructors ***/
 

diff --git a/numpy/core/src/multiarray/arrayobject.c b/numpy/core/src/multiarray/arrayobject.c
@@ -371,6 +371,53 @@ PyArray_TypeNumFromName(char *str)
     return NPY_NOTYPE;
 }
 
+/*NUMPY_API
+ *
+ * If UPDATEIFCOPY and self has data, reset the base WRITEABLE flag,
+ * copy the local data to base, release the local data, and set flags
+ * appropriately. Return 0 if not relevant, 1 if success, < 0 on failure
+ */
+NPY_NO_EXPORT int
+PyArray_ResolveUpdateIfCopy(PyArrayObject * self)
+{
+    PyArrayObject_fields *fa = (PyArrayObject_fields *)self;
+    if (fa && fa->base) {
+        if (fa->flags & NPY_ARRAY_UPDATEIFCOPY) {
+            /*
+             * UPDATEIFCOPY means that fa->base's data
+             * should be updated with the contents
+             * of self.
+             * fa->base->flags is not WRITEABLE to protect the relationship
+             * unlock it.
+             */
+            int retval = 0;
+            PyArray_ENABLEFLAGS(((PyArrayObject *)fa->base),
+                                                    NPY_ARRAY_WRITEABLE);
+            PyArray_CLEARFLAGS(self, NPY_ARRAY_UPDATEIFCOPY);
+            retval = PyArray_CopyAnyInto((PyArrayObject *)fa->base, self);
+            if (retval < 0) {
+                /* this should never happen, how did the two copies of data
+                 * get out of sync?
+                 */
+                return retval;
+            }
+            if ((fa->flags & NPY_ARRAY_OWNDATA) && fa->data) {
+                /* Free internal references if an Object array */
+                if (PyDataType_FLAGCHK(fa->descr, NPY_ITEM_REFCOUNT)) {
+                    Py_INCREF(self); /*hold on to self */
+                    PyArray_XDECREF(self);
+                    Py_DECREF(self); 
+                }
+                npy_free_cache(fa->data, PyArray_NBYTES(self));
+                fa->data = NULL;
+                PyArray_CLEARFLAGS(self, NPY_ARRAY_OWNDATA);
+            }
+            return 1;
+        }
+    }
+    return 0;
+ }
+
 /*********************** end C-API functions **********************/
 
 /* array object functions */
@@ -386,27 +433,21 @@ array_dealloc(PyArrayObject *self)
         PyObject_ClearWeakRefs((PyObject *)self);
     }
     if (fa->base) {
-        /*
-         * UPDATEIFCOPY means that base points to an
-         * array that should be updated with the contents
-         * of this array upon destruction.
-         * fa->base->flags must have been WRITEABLE
-         * (checked previously) and it was locked here
-         * thus, unlock it.
-         */
-        if (fa->flags & NPY_ARRAY_UPDATEIFCOPY) {
-            PyArray_ENABLEFLAGS(((PyArrayObject *)fa->base),
-                                                    NPY_ARRAY_WRITEABLE);
-            Py_INCREF(self); /* hold on to self in next call */
-            if (PyArray_CopyAnyInto((PyArrayObject *)fa->base, self) < 0) {
-                PyErr_Print();
-                PyErr_Clear();
-            }
-            /*
-             * Don't need to DECREF -- because we are deleting
-             *self already...
-             */
+        int retval;
+        Py_INCREF(self); /* hold on to self in next call */
+        retval = PyArray_ResolveUpdateIfCopy(self);
+        if (retval < 0)
+        {
+            PyErr_Print();
+            PyErr_Clear();
+        }
+#if defined(PPY_VERSION) || defined(DEPRECATE_UPDATIFCOPY)
+        if (retval > 0) {
+            DEPRECATE("UPDATEIFCOPY resolution in array_dealloc is "
+                      "incompatible with PyPy and will be removed in "
+                      "the future");
         }
+#endif
         /*
          * In any case base is pointing to something that we need
          * to DECREF -- either a view or a buffer object
@@ -426,7 +467,6 @@ array_dealloc(PyArrayObject *self)
         }
         npy_free_cache(fa->data, PyArray_NBYTES(self));
     }
-
     /* must match allocation in PyArray_NewFromDescr */
     npy_free_cache_dim(fa->dimensions, 2 * fa->nd);
     Py_DECREF(fa->descr);

diff --git a/numpy/core/src/multiarray/calculation.c b/numpy/core/src/multiarray/calculation.c
@@ -136,6 +136,7 @@ PyArray_ArgMax(PyArrayObject *op, int axis, PyArrayObject *out)
     Py_DECREF(ap);
     /* Trigger the UPDATEIFCOPY if necessary */
     if (out != NULL && out != rp) {
+        PyArray_ResolveUpdateIfCopy(rp);
         Py_DECREF(rp);
         rp = out;
         Py_INCREF(rp);
@@ -251,6 +252,7 @@ PyArray_ArgMin(PyArrayObject *op, int axis, PyArrayObject *out)
     Py_DECREF(ap);
     /* Trigger the UPDATEIFCOPY if necessary */
     if (out != NULL && out != rp) {
+        PyArray_ResolveUpdateIfCopy(rp);
         Py_DECREF(rp);
         rp = out;
         Py_INCREF(rp);
@@ -1153,6 +1155,7 @@ PyArray_Clip(PyArrayObject *self, PyObject *min, PyObject *max, PyArrayObject *o
     Py_XDECREF(maxa);
     Py_DECREF(newin);
     /* Copy back into out if out was not already a nice array. */
+    PyArray_ResolveUpdateIfCopy(newout);
     Py_DECREF(newout);
     return (PyObject *)out;
 

diff --git a/numpy/core/src/multiarray/cblasfuncs.c b/numpy/core/src/multiarray/cblasfuncs.c
@@ -770,6 +770,7 @@ cblas_matrixproduct(int typenum, PyArrayObject *ap1, PyArrayObject *ap2,
     Py_DECREF(ap2);
 
     /* Trigger possible copyback into `result` */
+    PyArray_ResolveUpdateIfCopy(out_buf);
     Py_DECREF(out_buf);
 
     return PyArray_Return(result);

diff --git a/numpy/core/src/multiarray/compiled_base.c b/numpy/core/src/multiarray/compiled_base.c
@@ -401,6 +401,7 @@ arr_insert(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwdict)
 
     Py_XDECREF(values);
     Py_XDECREF(mask);
+    PyArray_ResolveUpdateIfCopy(array);
     Py_DECREF(array);
     Py_RETURN_NONE;
 

diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
@@ -1801,6 +1801,7 @@ PyArray_FromAny(PyObject *op, PyArray_Descr *newtype, int min_depth,
         }
         else {
             ret = (PyArrayObject *)PyArray_FromArray(arr, newtype, flags);
+            PyArray_ResolveUpdateIfCopy(arr); /* prevent spurious warnings */
             Py_DECREF(arr);
         }
     }

diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c
@@ -234,6 +234,7 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis,
     Py_XDECREF(self);
     if (out != NULL && out != obj) {
         Py_INCREF(out);
+        PyArray_ResolveUpdateIfCopy(obj);
         Py_DECREF(obj);
         obj = out;
     }
@@ -406,6 +407,7 @@ PyArray_PutTo(PyArrayObject *self, PyObject* values0, PyObject *indices0,
     Py_XDECREF(values);
     Py_XDECREF(indices);
     if (copied) {
+        PyArray_ResolveUpdateIfCopy(self);
         Py_DECREF(self);
     }
     Py_RETURN_NONE;
@@ -523,6 +525,7 @@ PyArray_PutMask(PyArrayObject *self, PyObject* values0, PyObject* mask0)
     Py_XDECREF(values);
     Py_XDECREF(mask);
     if (copied) {
+        PyArray_ResolveUpdateIfCopy(self);
         Py_DECREF(self);
     }
     Py_RETURN_NONE;
@@ -768,6 +771,7 @@ PyArray_Choose(PyArrayObject *ip, PyObject *op, PyArrayObject *out,
     PyDataMem_FREE(mps);
     if (out != NULL && out != obj) {
         Py_INCREF(out);
+        PyArray_ResolveUpdateIfCopy(obj);
         Py_DECREF(obj);
         obj = out;
     }

diff --git a/numpy/core/src/multiarray/iterators.c b/numpy/core/src/multiarray/iterators.c
@@ -1151,6 +1151,7 @@ iter_richcompare(PyArrayIterObject *self, PyObject *other, int cmp_op)
         return NULL;
     }
     ret = array_richcompare(new, other, cmp_op);
+    PyArray_ResolveUpdateIfCopy(new);
     Py_DECREF(new);
     return ret;
 }

diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
@@ -1062,6 +1062,7 @@ PyArray_MatrixProduct2(PyObject *op1, PyObject *op2, PyArrayObject* out)
     Py_DECREF(ap2);
 
     /* Trigger possible copy-back into `result` */
+    PyArray_ResolveUpdateIfCopy(out_buf);
     Py_DECREF(out_buf);
 
     return (PyObject *)result;

diff --git a/numpy/core/src/multiarray/nditer_constr.c b/numpy/core/src/multiarray/nditer_constr.c
@@ -2927,6 +2927,7 @@ npyiter_allocate_arrays(NpyIter *iter,
                     Py_DECREF(temp);
                     return 0;
                 }
+                op_flags[iop] |= NPY_ITER_USE_CONTEXT_MANAGER;
             }
 
             Py_DECREF(op[iop]);

diff --git a/numpy/core/src/multiarray/nditer_pywrap.c b/numpy/core/src/multiarray/nditer_pywrap.c
@@ -24,6 +24,13 @@ struct NewNpyArrayIterObject_tag {
     NpyIter *iter;
     /* Flag indicating iteration started/stopped */
     char started, finished;
+    /* inside context manager
+       0 == NO, not needed 
+       1 == YES
+       2 == needed, (UPDATEIFCOPY flag is true). If this value is
+            reached and nditer API function called, fail on PyPy
+       3 == YES, and needed, all is good*/
+    char managed;
     /* Child to update for nested iteration */
     NewNpyArrayIterObject *nested_child;
     /* Cached values from the iterator */
@@ -710,6 +717,12 @@ npyiter_convert_ops(PyObject *op_in, PyObject *op_flags_in,
                 *nop_out = 0;
                 return 0;
             }
+            {
+                PyArrayObject_fields * fa = (PyArrayObject_fields*)ao;
+                if (fa->base && (fa->flags & NPY_ARRAY_UPDATEIFCOPY)) {
+                    op_flags[iop] |= NPY_ITER_USE_CONTEXT_MANAGER;
+                }
+            }
             Py_DECREF(op[iop]);
             op[iop] = ao;
         }
@@ -771,6 +784,13 @@ npyiter_init(NewNpyArrayIterObject *self, PyObject *args, PyObject *kwds)
         goto fail;
     }
 
+    self->managed = 0;
+    for (iop = 0; iop < nop; ++iop) {
+        if (op_flags[iop] & NPY_ITER_USE_CONTEXT_MANAGER)
+            self->managed |= 2; /* context manager recommended/required */
+        op_flags[iop] &= ~NPY_ITER_USE_CONTEXT_MANAGER;
+    }
+
     /* op_request_dtypes */
     if (op_dtypes_in != NULL && op_dtypes_in != Py_None &&
             npyiter_convert_dtypes(op_dtypes_in,
@@ -1063,6 +1083,14 @@ NpyIter_NestedIters(PyObject *NPY_UNUSED(self),
             Py_DECREF(ret);
             goto fail;
         }
+#ifdef PYPY_VERSION
+        if (iter->managed == 2)
+        {
+            PyErr_SetString(PyExc_ValueError,
+                    "'updateifcopy/copy' not supported on PyPy");
+            goto fail;
+        }
+#endif
 
         if (inest < nnest-1) {
             iter->iter = NpyIter_AdvancedNew(nop, op, flags, order,
@@ -1404,7 +1432,13 @@ static PyObject *npyiter_value_get(NewNpyArrayIterObject *self)
                 "Iterator is past the end");
         return NULL;
     }
-
+#ifdef PYPY_VERSION
+    if (self->managed == 2) {
+        PyErr_SetString(PyExc_ValueError,
+        "'updateifcopy', 'readwrite' flags must be used in a context manager");
+        return NULL;
+    }
+#endif
     nop = NpyIter_GetNOp(self->iter);
 
     /* Return an array  or tuple of arrays with the values */
@@ -1441,7 +1475,14 @@ static PyObject *npyiter_operands_get(NewNpyArrayIterObject *self)
                 "Iterator is invalid");
         return NULL;
     }
-
+#ifdef PYPY_VERSION
+    if (self->managed == 2) {
+        PyErr_SetString(PyExc_ValueError,
+        "'updateifcopy', 'readwrite' flags must be used in a context manager");
+        return NULL;
+    }
+#endif
+
     nop = NpyIter_GetNOp(self->iter);
     operands = self->operands;
 
@@ -2329,6 +2370,37 @@ npyiter_ass_subscript(NewNpyArrayIterObject *self, PyObject *op,
     return -1;
 }
 
+static PyObject *
+npyiter_self(NewNpyArrayIterObject *self)
+{
+    if (self->iter == NULL)
+    {
+        PyErr_SetString(PyExc_ValueError, "operation on  non-initialized iterator");
+        return NULL;
+    }
+    self->managed |= 1;
+    Py_INCREF(self);
+    return (PyObject *)self;
+}
+
+static PyObject *
+npyiter_exit(PyObject *self, PyObject *args)
+{
+    int iop, nop;
+    PyArrayObject **object;
+    NpyIter *iter = ((NewNpyArrayIterObject*)self)->iter;
+    if (iter == NULL)
+        Py_RETURN_NONE;
+    nop = NpyIter_GetNOp(iter);
+    object = NpyIter_GetOperandArray(iter);
+    for(iop = 0; iop < nop; ++iop, ++object) {
+        if (PyArray_ResolveUpdateIfCopy(*object) < 0)
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+
 static PyMethodDef npyiter_methods[] = {
     {"reset",
         (PyCFunction)npyiter_reset,
@@ -2354,6 +2426,10 @@ static PyMethodDef npyiter_methods[] = {
     {"debug_print",
         (PyCFunction)npyiter_debug_print,
         METH_NOARGS, NULL},
+    {"__enter__", (PyCFunction)npyiter_self,
+         METH_NOARGS,  NULL},
+    {"__exit__",  (PyCFunction)npyiter_exit,
+         METH_VARARGS, NULL},
     {NULL, NULL, 0, NULL},
 };
-Original file line number
+Diff line change
@@ Expand Up @@
             }
             else {
                 ret = (PyArrayObject *)PyArray_FromArray(arr, newtype, flags);
+                PyArray_ResolveUpdateIfCopy(arr); /* prevent spurious warnings */
                 Py_DECREF(arr);
             }
         }
@@ Expand Down @@