ENH: iter: Add check for buffered output allocation without delayed bufalloc

When allocating an output and specifying it's to be read, most iteration
modes allow the code to fill the allocated output with a value before
starting actual iteration.  When buffering is enabled and delayed
buffer allocation is not, this won't work, so throw an error

3 files changed, 52 insertions, 14 deletions

diff --git a/numpy/core/src/multiarray/new_iterator.c.src b/numpy/core/src/multiarray/new_iterator.c.src
index 25491166b..32ba8f4d3 100644
--- a/numpy/core/src/multiarray/new_iterator.c.src
+++ b/numpy/core/src/multiarray/new_iterator.c.src
@@ -232,6 +232,7 @@ npyiter_prepare_one_operand(PyArrayObject **op,
                         char **op_dataptr,
                         PyArray_Descr *op_request_dtype,
                         PyArray_Descr** op_dtype,
+                        npy_uint32 flags,
                         npy_uint32 op_flags, char *op_itflags);
 static int
 npyiter_prepare_operands(npy_intp niter, PyArrayObject **op_in,
@@ -239,6 +240,7 @@ npyiter_prepare_operands(npy_intp niter, PyArrayObject **op_in,
                     char **op_dataptr,
                     PyArray_Descr **op_request_dtypes,
                     PyArray_Descr **op_dtype,
+                    npy_uint32 flags,
                     npy_uint32 *op_flags, char *op_itflags);
 static int
 npyiter_check_casting(npy_intp niter, PyArrayObject **op,
@@ -379,6 +381,7 @@ NpyIter_MultiNew(npy_intp niter, PyArrayObject **op_in, npy_uint32 flags,
     /* Prepare all the operands */
     if (!npyiter_prepare_operands(niter, op_in, op, op_dataptr,
                         op_request_dtypes, op_dtype,
+                        flags,
                         op_flags, op_itflags)) {
         PyArray_free(iter);
         return NULL;
@@ -2351,6 +2354,7 @@ npyiter_prepare_one_operand(PyArrayObject **op,
                         char **op_dataptr,
                         PyArray_Descr *op_request_dtype,
                         PyArray_Descr **op_dtype,
+                        npy_uint32 flags,
                         npy_uint32 op_flags, char *op_itflags)
 {
     /* NULL operands must be automatically allocated outputs */
@@ -2358,24 +2362,32 @@ npyiter_prepare_one_operand(PyArrayObject **op,
         /* ALLOCATE should be enabled */
         if (!(op_flags&NPY_ITER_ALLOCATE)) {
             PyErr_SetString(PyExc_ValueError,
-                    "Iterator input was NULL, but automatic allocation as an "
+                    "Iterator operand was NULL, but automatic allocation as an "
                     "output wasn't requested");
             return 0;
         }
-        /* Reading should be disabled */
-        /* Disable this check for now, since except in one case (buffering
-         * without delayed buffer allocation), the caller can use
-         * NpyIter_GetObjectArray to initialize it after the iterator
-         * is created.
-         * TODO: Maybe reenable it when buffering is on, but delayed buffer
-         * allocation is off.
-        if ((*op_itflags)&NPY_OP_ITFLAG_READ) {
+        /* Writing should be enabled */
+        if (!((*op_itflags)&NPY_OP_ITFLAG_WRITE)) {
             PyErr_SetString(PyExc_ValueError,
                     "Automatic allocation was requested for an iterator "
-                    "operand, but it wasn't flagged as write only");
+                    "operand, but it wasn't flagged for writing");
+            return 0;
+        }
+        /*
+         * Reading should be disabled if buffering is enabled without
+         * also enabling NPY_ITER_DELAY_BUFALLOC.  In all other cases,
+         * the caller may initialize the allocated operand to a value
+         * before beginning iteration.
+         */
+        if (((flags&(NPY_ITER_BUFFERED|
+                        NPY_ITER_DELAY_BUFALLOC)) == NPY_ITER_BUFFERED) &&
+                ((*op_itflags)&NPY_OP_ITFLAG_READ)) {
+            PyErr_SetString(PyExc_ValueError,
+                    "Automatic allocation was requested for an iterator "
+                    "operand, and it was flagged as readable, but buffering "
+                    " without delayed allocation was enabled");
             return 0;
         }
-        */
         *op_dataptr = NULL;
         /* If a requested dtype was provided, use it, otherwise NULL */
         Py_XINCREF(op_request_dtype);
@@ -2513,6 +2525,7 @@ npyiter_prepare_operands(npy_intp niter, PyArrayObject **op_in,
                     char **op_dataptr,
                     PyArray_Descr **op_request_dtypes,
                     PyArray_Descr **op_dtype,
+                    npy_uint32 flags,
                     npy_uint32 *op_flags, char *op_itflags)
 {
     npy_intp iiter, i;
@@ -2539,6 +2552,7 @@ npyiter_prepare_operands(npy_intp niter, PyArrayObject **op_in,
                         &op_dataptr[iiter],
                         op_request_dtypes ? op_request_dtypes[iiter] : NULL,
                         &op_dtype[iiter],
+                        flags,
                         op_flags[iiter], &op_itflags[iiter])) {
             for (i = 0; i <= iiter; ++i) {
                 Py_XDECREF(op[i]);
diff --git a/numpy/core/src/multiarray/new_iterator_pywrap.c b/numpy/core/src/multiarray/new_iterator_pywrap.c
index b03ef4b14..3891bb806 100644
--- a/numpy/core/src/multiarray/new_iterator_pywrap.c
+++ b/numpy/core/src/multiarray/new_iterator_pywrap.c
@@ -1183,7 +1183,8 @@ npyiter_copy(NewNpyArrayIterObject *self)
 static PyObject *
 npyiter_iternext(NewNpyArrayIterObject *self)
 {
-    if (self->iter != NULL && !self->finished && self->iternext(self->iter)) {
+    if (self->iter != NULL && self->iternext != NULL &&
+                        !self->finished && self->iternext(self->iter)) {
         /* If there is nesting, the nested iterators should be reset */
         if (npyiter_resetbasepointers(self) != NPY_SUCCEED) {
             return NULL;
@@ -1279,7 +1280,12 @@ static PyObject *npyiter_value_get(NewNpyArrayIterObject *self)
             return NULL;
         }
         for (iiter = 0; iiter < niter; ++iiter) {
-            PyTuple_SET_ITEM(ret, iiter, npyiter_seq_item(self, iiter));
+            PyObject *a = npyiter_seq_item(self, iiter);
+            if (a == NULL) {
+                Py_DECREF(ret);
+                return NULL;
+            }
+            PyTuple_SET_ITEM(ret, iiter, a);
         }
     }
 
@@ -1350,7 +1356,7 @@ static PyObject *npyiter_itviews_get(NewNpyArrayIterObject *self)
 static PyObject *
 npyiter_next(NewNpyArrayIterObject *self)
 {
-    if (self->iter == NULL || self->finished) {
+    if (self->iter == NULL || self->iternext == NULL || self->finished) {
         return NULL;
     }
 
diff --git a/numpy/core/tests/test_new_iterator.py b/numpy/core/tests/test_new_iterator.py
index d012ed2f5..0f2879346 100644
--- a/numpy/core/tests/test_new_iterator.py
+++ b/numpy/core/tests/test_new_iterator.py
@@ -1073,6 +1073,18 @@ def test_iter_allocate_output_simple():
     assert_equal(i.operands[1].shape, a.shape)
     assert_equal(i.operands[1].dtype, np.dtype('f4'))
 
+def test_iter_allocate_output_buffered_readwrite():
+    # Allocated output with buffering + delay_bufalloc
+
+    a = arange(6)
+    i = newiter([a,None], ['buffered','delay_bufalloc'],
+                        [['readonly'],['allocate','readwrite']])
+    i.operands[1][:] = 1
+    i.reset()
+    for x in i:
+        x[1][()] += x[0][()]
+    assert_equal(i.operands[1], a+1)
+
 def test_iter_allocate_output_itorder():
     # The allocated output should match the iteration order
 
@@ -1182,6 +1194,12 @@ def test_iter_allocate_output_errors():
     a = arange(6)
     assert_raises(TypeError, newiter, [a,None], [],
                         [['writeonly'],['writeonly','allocate']])
+    # Allocated output should be flagged for writing
+    assert_raises(ValueError, newiter, [a,None], [],
+                        [['readonly'],['allocate','readonly']])
+    # Allocated output can't have buffering without delayed bufalloc
+    assert_raises(ValueError, newiter, [a,None], ['buffered'],
+                                            ['allocate','readwrite'])
     # Must specify at least one input
     assert_raises(ValueError, newiter, [None,None], [],
                         [['writeonly','allocate'],