[pypy-commit] pypy nditer-revisited: test, fix division by 0

[mattip]

Author: mattip <matti.picus at gmail.com>
Branch: nditer-revisited
Changeset: r78683:fb5def8bce10
Date: 2015-07-27 01:51 +0300
http://bitbucket.org/pypy/pypy/changeset/fb5def8bce10/

Log:	test, fix division by 0

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
@@ -303,6 +303,8 @@
         _shape = [shape[-1]]  + old_iter.slice_shape
         _backstride = [(shape[-1] - 1) * strides[-1]] + old_iter.slice_backstride
         fastest = shape[-1]
+    if fastest == 0:
+        return old_iter
     if flat:
         _shape = [support.product(_shape)]
         if len(_stride) > 1:
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
@@ -190,6 +190,17 @@
         assert a.shape == (2, 3, 4)
         assert (a == arange(24).reshape(2, 3, 4)).all()
 
+    def test_zerosize(self):
+        from numpy import nditer, array
+        for a in [ array([]), array([1]), array([1, 2]) ]:
+            buffersize = max(16 * 1024 ** 2 // a.itemsize, 1)
+            r = []
+            for chunk in nditer(a, 
+                    flags=['external_loop', 'buffered', 'zerosize_ok'],
+                    buffersize=buffersize, order='C'):
+                r.append(chunk)
+            assert (r == a).all()
+
     def test_op_dtype(self):
         from numpy import arange, nditer, sqrt, array
         a = arange(6).reshape(2,3) - 3