Chunkwise iteration over arrays.

Closes #398.

Author: jeromekelleher

zarr/core.py

@@ -424,6 +424,18 @@ def __array__(self, *args):
             a = a.astype(args[0])
         return a
 
+    def __iter__(self):
+        if len(self.shape) == 0:
+            # Same error as numpy
+            raise TypeError("iteration over a 0-d array")
+        # Avoid repeatedly decompressing chunks by iterating over the chunks
+        # in the first dimension.
+        chunk_size = self.chunks[0]
+        for j in range(self.shape[0]):
+            if j % chunk_size == 0:
+                chunk = self[j: j + chunk_size][:]
+            yield chunk[j % chunk_size]
+
     def __len__(self):
         if self.shape:
             return self.shape[0]

zarr/tests/test_core.py

@@ -7,6 +7,7 @@
 import pickle
 import os
 import warnings
+import itertools
 
 
 import numpy as np
@@ -1155,6 +1156,40 @@ def test_object_codec_warnings(self):
             # provide object_codec, but not object dtype
             self.create_array(shape=10, chunks=5, dtype='i4', object_codec=JSON())
 
+    def test_zero_d_iter(self):
+        a = np.array(1, dtype=int)
+        z = self.create_array(shape=a.shape, dtype=int)
+        z[...] = a
+        with pytest.raises(TypeError):
+            # noinspection PyStatementEffect
+            list(a)
+        with pytest.raises(TypeError):
+            # noinspection PyStatementEffect
+            list(z)
+
+    def test_iter(self):
+        params = (
+            ((1,), (1,)),
+            ((2,), (1,)),
+            ((1,), (2,)),
+            ((3,), (3,)),
+            ((1000,), (100,)),
+            ((100,), (1000,)),
+            ((1, 100), (1, 1)),
+            ((1, 0), (1, 1)),
+            ((0, 1), (1, 1)),
+            ((0, 1), (2, 1)),
+            ((100, 1), (3, 1)),
+            ((100, 100), (10, 10)),
+            ((10, 10, 10), (3, 3, 3)),
+        )
+        for shape, chunks in params:
+            z = self.create_array(shape=shape, chunks=chunks, dtype=int)
+            a = np.arange(np.product(shape)).reshape(shape)
+            z[:] = a
+            for expect, actual in itertools.zip_longest(a, z):
+                assert_array_equal(expect, actual)
+
 
 class TestArrayWithPath(TestArray):