[pypy-commit] pypy numpypy-nditer: cleanup, add failing test that proves refactoring is needed

[mattip]

Author: Matti Picus <matti.picus at gmail.com>
Branch: numpypy-nditer
Changeset: r70630:f0bb0496eeef
Date: 2014-04-14 15:56 +0300
http://bitbucket.org/pypy/pypy/changeset/f0bb0496eeef/

Log:	cleanup, add failing test that proves refactoring is needed

diff --git a/pypy/module/micronumpy/nditer.py b/pypy/module/micronumpy/nditer.py
--- a/pypy/module/micronumpy/nditer.py
+++ b/pypy/module/micronumpy/nditer.py
@@ -189,28 +189,6 @@
                 'Iterator flag EXTERNAL_LOOP cannot be used if an index or '
                 'multi-index is being tracked'))
 
-def get_iter(space, order, arr, shape):
-    imp = arr.implementation
-    if order == 'K' or (order == 'C' and imp.order == 'C'):
-        backward = False
-    elif order =='F' and imp.order == 'C':
-        backward = True
-    else:
-        raise OperationError(space.w_NotImplementedError, space.wrap(
-                'not implemented yet'))
-    if (imp.strides[0] < imp.strides[-1] and not backward) or \
-       (imp.strides[0] > imp.strides[-1] and backward):
-        # flip the strides. Is this always true for multidimension?
-        strides = [imp.strides[i] for i in range(len(imp.strides) - 1, -1, -1)]
-        backstrides = [imp.backstrides[i] for i in range(len(imp.backstrides) - 1, -1, -1)]
-        shape = [imp.shape[i] for i in range(len(imp.shape) - 1, -1, -1)]
-    else:
-        strides = imp.strides
-        backstrides = imp.backstrides
-    r = calculate_broadcast_strides(strides, backstrides, imp.shape,
-                                    shape, backward)
-    return ArrayIter(imp, imp.get_size(), shape, r[0], r[1])
-
 def is_backward(imp, order):
     if order == 'K' or (order == 'C' and imp.order == 'C'):
         return False
@@ -219,11 +197,28 @@
     else:
         raise NotImplementedError('not implemented yet')
 
+def get_iter(space, order, arr, shape):
+    imp = arr.implementation
+    backward = is_backward(imp, order)
+    if (imp.strides[0] < imp.strides[-1] and not backward) or \
+       (imp.strides[0] > imp.strides[-1] and backward):
+        # flip the strides. Is this always true for multidimension?
+        strides = imp.strides[:]
+        backstrides = imp.backstrides[:]
+        shape = imp.shape[:]
+        strides.reverse()
+        backstrides.reverse()
+        shape.reverse()
+    else:
+        strides = imp.strides
+        backstrides = imp.backstrides
+    r = calculate_broadcast_strides(strides, backstrides, imp.shape,
+                                    shape, backward)
+    return ArrayIter(imp, imp.get_size(), shape, r[0], r[1])
+
 def get_external_loop_iter(space, order, arr, shape):
     imp = arr.implementation
-
     backward = is_backward(imp, order)
-
     return SliceIterator(arr, imp.strides, imp.backstrides, shape, order=order, backward=backward)
 
 def convert_to_array_or_none(space, w_elem):
diff --git a/pypy/module/micronumpy/test/test_nditer.py b/pypy/module/micronumpy/test/test_nditer.py
--- a/pypy/module/micronumpy/test/test_nditer.py
+++ b/pypy/module/micronumpy/test/test_nditer.py
@@ -54,7 +54,14 @@
         assert (array(r) == [[ 0, 12], [ 4, 16], [ 8, 20], [ 1, 13], [ 5, 17], [ 9, 21], [ 2, 14], [ 6, 18], [10, 22], [ 3, 15], [ 7, 19], [11, 23]]).all()
         e = raises(ValueError, 'r[0][0] = 0')
         assert str(e.value) == 'assignment destination is read-only'
-           
+        r = []
+        for x in nditer(a.T, flags=['external_loop'], order='F'):
+            r.append(x)
+        array_r = array(r)
+        assert len(array_r.shape) == 2
+        assert array_r.shape == (1,24)
+        assert (array(r) == arange(24)).all()
+
     def test_index(self):
         from numpy import arange, nditer
         a = arange(6).reshape(2,3)