Add recurse keyword to array_keys and arrays (#458)

* Adds recurse keyword to array_keys and arrays

* Adds item to release notes

Author: jrbourbeau

docs/release.rst

@@ -4,6 +4,9 @@ Release notes
 Upcoming Release
 ----------------
 
+* Add ``recurse`` keyword to ``Group.array_keys`` and ``Group.arrays`` methods.
+  By :user:`James Bourbeau <jrbourbeau>`; :issue:`458`
+
 * Use uniform chunking for all dimensions when specifying ``chunks`` as an integer.
   Also adds support for specifying ``-1`` to chunk across an entire dimension.
   By :user:`James Bourbeau <jrbourbeau>`; :issue:`456`.

zarr/hierarchy.py

@@ -406,9 +406,16 @@ def groups(self):
                                  cache_attrs=self.attrs.cache,
                                  synchronizer=self._synchronizer)
 
-    def array_keys(self):
+    def array_keys(self, recurse=False):
         """Return an iterator over member names for arrays only.
 
+        Parameters
+        ----------
+        recurse : recurse, optional
+            Option to return member names for all arrays, even from groups
+            below the current one. If False, only member names for arrays in
+            the current group will be returned. Default value is False.
+
         Examples
         --------
         >>> import zarr
@@ -421,14 +428,20 @@ def array_keys(self):
         ['baz', 'quux']
 
         """
-        for key in sorted(listdir(self._store, self._path)):
-            path = self._key_prefix + key
-            if contains_array(self._store, path):
-                yield key
+        return self._array_iter(keys_only=True,
+                                method='array_keys',
+                                recurse=recurse)
 
-    def arrays(self):
+    def arrays(self, recurse=False):
         """Return an iterator over (name, value) pairs for arrays only.
 
+        Parameters
+        ----------
+        recurse : recurse, optional
+            Option to return (name, value) pairs for all arrays, even from groups
+            below the current one. If False, only (name, value) pairs for arrays in
+            the current group will be returned. Default value is False.
+
         Examples
         --------
         >>> import zarr
@@ -443,13 +456,19 @@ def arrays(self):
         quux <class 'zarr.core.Array'>
 
         """
+        return self._array_iter(keys_only=False,
+                                method='arrays',
+                                recurse=recurse)
+
+    def _array_iter(self, keys_only, method, recurse):
         for key in sorted(listdir(self._store, self._path)):
             path = self._key_prefix + key
             if contains_array(self._store, path):
-                yield key, Array(self._store, path=path, read_only=self._read_only,
-                                 chunk_store=self._chunk_store,
-                                 cache_attrs=self.attrs.cache,
-                                 synchronizer=self._synchronizer)
+                yield key if keys_only else (key, self[key])
+            elif recurse and contains_group(self._store, path):
+                group = self[key]
+                for i in getattr(group, method)(recurse=recurse):
+                    yield i
 
     def visitvalues(self, func):
         """Run ``func`` on each object.

zarr/tests/test_hierarchy.py

@@ -636,6 +636,30 @@ def test_empty_getitem_contains_iterators(self):
         assert 0 == len(g)
         assert 'foo' not in g
 
+    def test_iterators_recurse(self):
+        # setup
+        g1 = self.create_group()
+        g2 = g1.create_group('foo/bar')
+        d1 = g2.create_dataset('/a/b/c', shape=1000, chunks=100)
+        d1[:] = np.arange(1000)
+        d2 = g1.create_dataset('foo/baz', shape=3000, chunks=300)
+        d2[:] = np.arange(3000)
+        d3 = g2.create_dataset('zab', shape=2000, chunks=200)
+        d3[:] = np.arange(2000)
+
+        # test recursive array_keys
+        array_keys = list(g1['foo'].array_keys(recurse=False))
+        array_keys_recurse = list(g1['foo'].array_keys(recurse=True))
+        assert len(array_keys_recurse) > len(array_keys)
+        assert sorted(array_keys_recurse) == ['baz', 'zab']
+
+        # test recursive arrays
+        arrays = list(g1['foo'].arrays(recurse=False))
+        arrays_recurse = list(g1['foo'].arrays(recurse=True))
+        assert len(arrays_recurse) > len(arrays)
+        assert 'zab' == arrays_recurse[0][0]
+        assert g1['foo']['bar']['zab'] == arrays_recurse[0][1]
+
     def test_getattr(self):
         # setup
         g1 = self.create_group()