Merge pull request #32 from braingram/to_internal_tests

To internal tests

Author: braingram

asdf_zarr/converter.py

@@ -32,14 +32,14 @@ def to_yaml_tree(self, obj, tag, ctx):
             chunk_key_block_index_map = {}
             for chunk_key in storage._iter_chunk_keys(obj, only_initialized=True):
                 data_callback = storage._generate_chunk_data_callback(obj, chunk_key)
-                asdf_key = chunk_store._chunk_asdf_keys.get(chunk_key, asdf.util.BlockKey())
-                block_index = ctx.find_block_index(asdf_key, data_callback)
+                asdf_key = chunk_store._chunk_asdf_keys.get(chunk_key, ctx.generate_block_key())
+                block_index = ctx.find_available_block_index(data_callback, asdf_key)
                 chunk_key_block_index_map[chunk_key] = block_index
             asdf_key = chunk_store._chunk_block_map_asdf_key
             if asdf_key is None:
-                asdf_key = asdf.util.BlockKey()
-            obj_dict["chunk_block_map"] = ctx.find_block_index(
-                asdf_key, storage._generate_chunk_map_callback(obj, chunk_key_block_index_map)
+                asdf_key = ctx.generate_block_key()
+            obj_dict["chunk_block_map"] = ctx.find_available_block_index(
+                storage._generate_chunk_map_callback(obj, chunk_key_block_index_map), asdf_key
             )
             return obj_dict
 
@@ -75,7 +75,7 @@ def from_yaml_tree(self, node, tag, ctx):
             return obj
 
         chunk_store = util.decode_storage(node["store"])
-        if "meta" in node:
+        if "meta_store" in node:
             # separate meta and chunk stores
             store = util.decode_storage(node["meta_store"])
         else:

asdf_zarr/storage.py

@@ -2,7 +2,6 @@
 import json
 import math
 
-import asdf
 import numpy
 import zarr
 
@@ -61,7 +60,7 @@ def to_internal(zarray):
         return zarray
     # make a new internal store based off an existing store
     internal_store = ConvertedInternalStore(zarray.chunk_store or zarray.store)
-    return zarr.open(internal_store)
+    return zarr.open(zarray.store, chunk_store=internal_store)
 
 
 class InternalStore(zarr.storage.Store):
@@ -139,11 +138,10 @@ def __init__(self, ctx, chunk_block_map_index, zarray_meta):
         # split into 4 chunks) the chunk_block_map will be
         # 4 x 5
         cdata_shape = tuple(math.ceil(s / c) for s, c in zip(zarray_meta["shape"], zarray_meta["chunks"]))
+        self._chunk_block_map_asdf_key = ctx.generate_block_key()
         self._chunk_block_map = numpy.frombuffer(
-            ctx.get_block_data_callback(chunk_block_map_index)(), dtype="int32"
+            ctx.get_block_data_callback(chunk_block_map_index, self._chunk_block_map_asdf_key)(), dtype="int32"
         ).reshape(cdata_shape)
-        self._chunk_block_map_asdf_key = asdf.util.BlockKey()
-        ctx.assign_block_key(chunk_block_map_index, self._chunk_block_map_asdf_key)
 
         self._chunk_block_map_asdf_key = None
 
@@ -154,10 +152,9 @@ def __init__(self, ctx, chunk_block_map_index, zarray_meta):
             coord = tuple(coord)
             block_index = int(self._chunk_block_map[coord])
             chunk_key = self._sep.join((str(c) for c in tuple(coord)))
-            asdf_key = asdf.util.BlockKey()
+            asdf_key = ctx.generate_block_key()
             self._chunk_asdf_keys[chunk_key] = asdf_key
-            ctx.assign_block_key(block_index, asdf_key)
-            self._chunk_callbacks[chunk_key] = ctx.get_block_data_callback(block_index)
+            self._chunk_callbacks[chunk_key] = ctx.get_block_data_callback(block_index, asdf_key)
 
     def _sep_key(self, key):
         if self._sep is None:

asdf_zarr/tests/test_zarr.py

@@ -2,18 +2,22 @@
 import itertools
 
 import asdf
+import asdf_zarr
+import asdf_zarr.storage
 import numpy
 import pytest
 import zarr
-from zarr.storage import DirectoryStore, NestedDirectoryStore
+from zarr.storage import DirectoryStore, KVStore, MemoryStore, NestedDirectoryStore, TempStore
 
 
-def create_zarray(shape=None, chunks=None, dtype="f8", store=None):
+def create_zarray(shape=None, chunks=None, dtype="f8", store=None, chunk_store=None):
     if shape is None:
         shape = (6, 9)
     if chunks is None:
         chunks = [max(1, d // 3) for d in shape]
-    arr = zarr.creation.create((6, 9), store=store, chunks=chunks, dtype=dtype, compressor=None)
+    arr = zarr.creation.create(
+        (6, 9), store=store, chunk_store=chunk_store, chunks=chunks, dtype=dtype, compressor=None
+    )
     for chunk_index in itertools.product(*[range(c) for c in arr.cdata_shape]):
         inds = []
         for i, c in zip(chunk_index, arr.chunks):
@@ -25,14 +29,37 @@ def create_zarray(shape=None, chunks=None, dtype="f8", store=None):
 @pytest.mark.parametrize("copy_arrays", [True, False])
 @pytest.mark.parametrize("lazy_load", [True, False])
 @pytest.mark.parametrize("compression", ["input", "zlib"])
-@pytest.mark.parametrize("store_type", [DirectoryStore, NestedDirectoryStore])
-def test_write_to(tmp_path, copy_arrays, lazy_load, compression, store_type):
-    store1 = store_type(tmp_path / "zarr_array_1")
-    store2 = store_type(tmp_path / "zarr_array_2")
+@pytest.mark.parametrize("store_type", [DirectoryStore, KVStore, MemoryStore, NestedDirectoryStore, TempStore])
+@pytest.mark.parametrize("to_internal", [True, False])
+@pytest.mark.parametrize("meta_store", [True, False])
+def test_write_to(tmp_path, copy_arrays, lazy_load, compression, store_type, to_internal, meta_store):
+    if store_type in (DirectoryStore, NestedDirectoryStore):
+        store1 = store_type(tmp_path / "zarr_array_1")
+        store2 = store_type(tmp_path / "zarr_array_2")
+    elif store_type is KVStore:
+        store1 = store_type({})
+        store2 = store_type({})
+    else:
+        store1 = store_type()
+        store2 = store_type()
+
+    # should meta be in a different store?
+    if meta_store:
+        chunk_store1 = store1
+        store1 = KVStore({})
+        chunk_store2 = store2
+        store2 = KVStore({})
+    else:
+        chunk_store1 = None
+        chunk_store2 = None
+
+    arr1 = create_zarray(store=store1, chunk_store=chunk_store1)
+    arr2 = create_zarray(store=store2, chunk_store=chunk_store2)
 
-    arr1 = create_zarray(store=store1)
-    arr2 = create_zarray(store=store2)
     arr2[:] = arr2[:] * -2
+    if to_internal:
+        arr1 = asdf_zarr.storage.to_internal(arr1)
+        arr2 = asdf_zarr.storage.to_internal(arr2)
     tree = {"arr1": arr1, "arr2": arr2}
 
     fn = tmp_path / "test.asdf"
@@ -42,8 +69,10 @@ def test_write_to(tmp_path, copy_arrays, lazy_load, compression, store_type):
     with asdf.open(fn, mode="r", copy_arrays=copy_arrays, lazy_load=lazy_load) as af:
         for n, a in (("arr1", arr1), ("arr2", arr2)):
             assert isinstance(af[n], zarr.core.Array)
-            # for these tests, data should not be converted to a different storage format
-            assert isinstance(af[n].chunk_store, store_type)
+            if to_internal or store_type in (KVStore, MemoryStore, TempStore):
+                assert isinstance(af[n].chunk_store, asdf_zarr.storage.InternalStore)
+            else:
+                assert isinstance(af[n].chunk_store, store_type)
             assert numpy.allclose(af[n], a)
 
 

asdf_zarr/util.py

@@ -2,7 +2,7 @@
 
 import fsspec
 import zarr.storage
-from zarr.storage import DirectoryStore, FSStore, NestedDirectoryStore, TempStore
+from zarr.storage import DirectoryStore, FSStore, KVStore, NestedDirectoryStore, TempStore
 
 
 def encode_storage(store):
@@ -33,6 +33,8 @@ def encode_storage(store):
         obj_dict["mode"] = store.mode
         # store.fs.to_json to get full filesystem (see fsspec.AbstractFileSystem.from_json)
         obj_dict["fs"] = store.fs.to_json()
+    elif isinstance(store, KVStore):
+        obj_dict["map"] = {k: store[k] for k in store}
     else:
         raise NotImplementedError(f"zarr.storage.Store subclass {store.__class__} not supported")
     return obj_dict
@@ -59,4 +61,6 @@ def decode_storage(obj_dict):  # TODO needs kwargs for dimension sep?
     if "fs" in kwargs and type_string == "FSStore":
         kwargs["fs"] = fsspec.AbstractFileSystem.from_json(kwargs["fs"])
         args.append(kwargs.pop("path"))
+    elif "map" in kwargs and type_string == "KVStore":
+        args.append(kwargs.pop("map"))
     return getattr(zarr.storage, type_string)(*args, **kwargs)