Allow chunk_store argument in open_consolidated() (#557)

Author: ajelenak

zarr/convenience.py

@@ -1179,4 +1179,5 @@ def open_consolidated(store, metadata_key='.zmetadata', mode='r+', **kwargs):
     meta_store = ConsolidatedMetadataStore(store, metadata_key=metadata_key)
 
     # pass through
-    return open(store=meta_store, chunk_store=store, mode=mode, **kwargs)
+    chunk_store = kwargs.pop('chunk_store', None) or store
+    return open(store=meta_store, chunk_store=chunk_store, mode=mode, **kwargs)

zarr/tests/test_convenience.py

@@ -165,6 +165,56 @@ def test_consolidate_metadata():
     open_consolidated(store, cache_attrs=True, synchronizer=None)
 
 
+def test_consolidated_with_chunk_store():
+    # setup initial data
+    store = MemoryStore()
+    chunk_store = MemoryStore()
+    z = group(store, chunk_store=chunk_store)
+    z.create_group('g1')
+    g2 = z.create_group('g2')
+    g2.attrs['hello'] = 'world'
+    arr = g2.create_dataset('arr', shape=(20, 20), chunks=(5, 5), dtype='f8')
+    assert 16 == arr.nchunks
+    assert 0 == arr.nchunks_initialized
+    arr.attrs['data'] = 1
+    arr[:] = 1.0
+    assert 16 == arr.nchunks_initialized
+
+    # perform consolidation
+    out = consolidate_metadata(store)
+    assert isinstance(out, Group)
+    assert '.zmetadata' in store
+    for key in ['.zgroup',
+                'g1/.zgroup',
+                'g2/.zgroup',
+                'g2/.zattrs',
+                'g2/arr/.zarray',
+                'g2/arr/.zattrs']:
+        del store[key]
+    # open consolidated
+    z2 = open_consolidated(store, chunk_store=chunk_store)
+    assert ['g1', 'g2'] == list(z2)
+    assert 'world' == z2.g2.attrs['hello']
+    assert 1 == z2.g2.arr.attrs['data']
+    assert (z2.g2.arr[:] == 1.0).all()
+    assert 16 == z2.g2.arr.nchunks
+    assert 16 == z2.g2.arr.nchunks_initialized
+
+    # test the data are writeable
+    z2.g2.arr[:] = 2
+    assert (z2.g2.arr[:] == 2).all()
+
+    # test invalid modes
+    with pytest.raises(ValueError):
+        open_consolidated(store, mode='a', chunk_store=chunk_store)
+    with pytest.raises(ValueError):
+        open_consolidated(store, mode='w', chunk_store=chunk_store)
+
+    # make sure keyword arguments are passed through without error
+    open_consolidated(store, cache_attrs=True, synchronizer=None,
+                      chunk_store=chunk_store)
+
+
 class TestCopyStore(unittest.TestCase):
 
     def setUp(self):