ENH: iter: Add simple tests for buffering

3 files changed, 106 insertions, 37 deletions

diff --git a/numpy/core/src/multiarray/lowlevel_strided_loops.c.src b/numpy/core/src/multiarray/lowlevel_strided_loops.c.src
index b51f1e680..8c62baace 100644
--- a/numpy/core/src/multiarray/lowlevel_strided_loops.c.src
+++ b/numpy/core/src/multiarray/lowlevel_strided_loops.c.src
@@ -766,7 +766,7 @@ PyArray_GetTransferFunction(int aligned,
         /* TODO wrap the cast in an alignment operation */
     }
 
-    /* TODO: write more complicated transfer code! */
+    /* TODO: write the more complicated transfer code! */
     *outstransfer = NULL;
     *outtransferdata = NULL;
     PyErr_SetString(PyExc_RuntimeError,
diff --git a/numpy/core/src/multiarray/new_iterator_pywrap.c b/numpy/core/src/multiarray/new_iterator_pywrap.c
index 0695d7f4c..1f9cd0b53 100644
--- a/numpy/core/src/multiarray/new_iterator_pywrap.c
+++ b/numpy/core/src/multiarray/new_iterator_pywrap.c
@@ -17,8 +17,8 @@ struct NewNpyArrayIterObject_tag {
     PyObject_HEAD
     /* The iterator */
     NpyIter *iter;
-    /* Flag indicating iteration stopped */
-    char finished;
+    /* Flag indicating iteration started/stopped */
+    char started, finished;
     /* Cached values from the iterator */
     NpyIter_IterNext_Fn iternext;
     NpyIter_GetCoords_Fn getcoords;
@@ -549,6 +549,9 @@ npyiter_init(NewNpyArrayIterObject *self, PyObject *args, PyObject *kwds)
     /* Cache some values for the member functions to use */
     npyiter_cache_values(self);
 
+    self->started = 0;
+    self->finished = 0;
+
     /* Release the references we got to the ops and dtypes */
     for (iiter = 0; iiter < niter; ++iiter) {
         Py_XDECREF(op[iiter]);
@@ -585,6 +588,7 @@ npyiter_reset(NewNpyArrayIterObject *self)
     }
 
     NpyIter_Reset(self->iter);
+    self->started = 0;
     self->finished = 0;
 
     Py_RETURN_NONE;
@@ -615,6 +619,7 @@ npyiter_remove_coords(NewNpyArrayIterObject *self)
     /* RemoveCoords invalidates cached values */
     npyiter_cache_values(self);
     /* RemoveCoords also resets the iterator */
+    self->started = 0;
     self->finished = 0;
 
     Py_RETURN_NONE;
@@ -633,6 +638,7 @@ npyiter_remove_inner_loop(NewNpyArrayIterObject *self)
     /* RemoveInnerLoop invalidates cached values */
     npyiter_cache_values(self);
     /* RemoveInnerLoop also resets the iterator */
+    self->started = 0;
     self->finished = 0;
 
     Py_RETURN_NONE;
@@ -765,22 +771,23 @@ static PyObject *npyiter_itviews_get(NewNpyArrayIterObject *self)
 static PyObject *
 npyiter_next(NewNpyArrayIterObject *self)
 {
-    PyObject *ret;
-
     if (self->iter == NULL || self->finished) {
         return NULL;
     }
 
-    ret = npyiter_value_get(self);
-    if (ret == NULL) {
-        return NULL;
+    /*
+     * Use the started flag for the Python iteration protocol to work
+     * when buffering is enabled.
+     */
+    if (self->started) {
+        if (!self->iternext(self->iter)) {
+            self->finished = 1;
+            return NULL;
+        }
     }
+    self->started = 1;
 
-    if (!self->iternext(self->iter)) {
-        self->finished = 1;
-    }
-    
-    return ret;
+    return npyiter_value_get(self);
 };
 
 static PyObject *npyiter_shape_get(NewNpyArrayIterObject *self)
@@ -875,6 +882,7 @@ static int npyiter_coords_set(NewNpyArrayIterObject *self, PyObject *value)
         if (NpyIter_GotoCoords(self->iter, coords) != NPY_SUCCEED) {
             return -1;
         }
+        self->started = 0;
         self->finished = 0;
         
         return 0;
@@ -928,6 +936,7 @@ static int npyiter_index_set(NewNpyArrayIterObject *self, PyObject *value)
         if (NpyIter_GotoIndex(self->iter, index) != NPY_SUCCEED) {
             return -1;
         }
+        self->started = 0;
         self->finished = 0;
         return 0;
     }
@@ -1123,10 +1132,11 @@ NPY_NO_EXPORT int
 npyiter_seq_ass_item(NewNpyArrayIterObject *self, Py_ssize_t i, PyObject *v)
 {
 
-    npy_intp niter;
+    npy_intp niter, innerloopsize, innerstride;
     char *dataptr;
     PyArray_Descr *dtype;
-    PyArrayObject *object;
+    PyArrayObject *tmp;
+    int ret;
 
     if (v == NULL) {
         PyErr_SetString(PyExc_ValueError,
@@ -1152,31 +1162,29 @@ npyiter_seq_ass_item(NewNpyArrayIterObject *self, Py_ssize_t i, PyObject *v)
 
     dataptr = self->dataptrs[i];
     dtype = self->dtypes[i];
-    object = self->objects[i];
 
     if (NpyIter_HasInnerLoop(self->iter)) {
-        /*
-         * TODO: When buffering is enabled for an operand, the object won't
-         *       correspond to the data, so that will have to be accounted for
-         */
-        return dtype->f->setitem(v, dataptr, object);
-    } else {
-        PyArrayObject *tmp;
-        int ret;
-        Py_INCREF(dtype);
-        /* TODO - there should be a better way than this... */
-        tmp = (PyArrayObject *)PyArray_NewFromDescr(&PyArray_Type, dtype,
-                                    1, self->innerloopsizeptr,
-                                    &self->innerstrides[i], dataptr,
-                                    NPY_WRITEABLE, NULL);
-        if (tmp == NULL) {
-            return -1;
-        }
-        PyArray_UpdateFlags(tmp, NPY_UPDATE_ALL);
-        ret = PyArray_CopyObject(tmp, v);
-        Py_DECREF(tmp);
-        return ret;
+        innerloopsize = 1;
+        innerstride = 0;
     }
+    else {
+        innerloopsize = *self->innerloopsizeptr;
+        innerstride = self->innerstrides[i];
+    }
+
+    /* TODO - there should be a better way than this... */
+    Py_INCREF(dtype);
+    tmp = (PyArrayObject *)PyArray_NewFromDescr(&PyArray_Type, dtype,
+                                1, &innerloopsize,
+                                &innerstride, dataptr,
+                                NPY_WRITEABLE, NULL);
+    if (tmp == NULL) {
+        return -1;
+    }
+    PyArray_UpdateFlags(tmp, NPY_UPDATE_ALL);
+    ret = PyArray_CopyObject(tmp, v);
+    Py_DECREF(tmp);
+    return ret;
 }
 
 static PyMethodDef npyiter_methods[] = {
diff --git a/numpy/core/tests/test_new_iterator.py b/numpy/core/tests/test_new_iterator.py
index 9002d231e..7054f0a47 100644
--- a/numpy/core/tests/test_new_iterator.py
+++ b/numpy/core/tests/test_new_iterator.py
@@ -1091,5 +1091,66 @@ def test_iter_remove_coords_inner_loop():
     assert_equal(i.itersize, 1)
     assert_equal(i.value, arange(24))
 
+def test_iter_buffering():
+    # Test buffering with several buffer sizes and types
+    arrays = []
+    # F-order swapped array
+    arrays.append(np.arange(24,
+                    dtype='c16').reshape(2,3,4).T.newbyteorder().byteswap())
+    # Contiguous 1-dimensional array
+    arrays.append(np.arange(10, dtype='f4'))
+    # Unaligned array
+    a = np.zeros((4*16+1,), dtype='i1')[1:]
+    a.dtype = 'i4'
+    a[:] = np.arange(16,dtype='i4')
+    arrays.append(a)
+    for a in arrays:
+        for buffersize in (1,2,3,5,8,11,16,1024):
+            vals = []
+            i = np.newiter(a, ['buffered','no_inner_iteration','force_c_order'],
+                           [['readonly','nbo_aligned']],
+                           buffersize=buffersize)
+            while not i.finished:
+                assert_(i[0].size <= buffersize)
+                vals.append(i[0].copy())
+                i.iternext()
+            assert_equal(np.concatenate(vals), a.ravel(order='C'))
+
+def test_iter_write_buffering():
+    # Test that buffering of writes is working
+
+    # F-order swapped array
+    a = np.arange(24).reshape(2,3,4).T.newbyteorder().byteswap()
+    i = np.newiter(a, ['buffered','force_c_order'],
+                   [['readwrite','nbo_aligned']],
+                   buffersize=16)
+    x = 0
+    while not i.finished:
+        i[0] = x
+        x += 1
+        i.iternext()
+    assert_equal(a.ravel(order='C'), np.arange(24))
+
+def test_iter_cast_buffering():
+    # Test that buffering can handle a simple cast
+
+    a = np.arange(10, dtype='f4')
+    i = np.newiter(a, ['buffered','no_inner_iteration'],
+                   [['readwrite','nbo_aligned','same_kind_casts']],
+                   op_dtypes=[np.dtype('f8')],
+                   buffersize=3)
+    for v in i:
+        print v
+        v[()] *= 2
+    
+    assert_equal(a, 2*np.arange(10, dtype='f4'))
+
+def test_iter_buffering_growinner():
+    # Test that the inner loop grows when no buffering is needed
+    a = np.arange(30)
+    i = np.newiter(a, ['buffered_growinner','no_inner_iteration'], buffersize=5)
+    # Should end up with just one inner loop here
+    assert_equal(i[0].size, a.size)
+
 if __name__ == "__main__":
     run_module_suite()