diff --git a/docs/release.rst b/docs/release.rst index 8626ace746..4d06185a36 100644 --- a/docs/release.rst +++ b/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 `; :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 `; :issue:`456` diff --git a/zarr/hierarchy.py b/zarr/hierarchy.py index d0ea600010..4659e75601 100644 --- a/zarr/hierarchy.py +++ b/zarr/hierarchy.py @@ -410,9 +410,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 @@ -425,14 +432,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 @@ -447,13 +460,19 @@ def arrays(self): quux """ + 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. diff --git a/zarr/tests/test_hierarchy.py b/zarr/tests/test_hierarchy.py index 30bcd484cf..52b8751db0 100644 --- a/zarr/tests/test_hierarchy.py +++ b/zarr/tests/test_hierarchy.py @@ -642,6 +642,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()