v3 compatibility

asdf_zarr/_zarr_compat.py

@@ -0,0 +1,12 @@
+import importlib
+
+# TODO: v3 is the default branch yet the version still reports as 2.x
+# for now, look for the "v2" submodule
+_is_v3 = importlib.util.find_spec("zarr.v2") is not None
+
+if _is_v3:
+    import zarr.v2 as zarr
+    import zarr.v2.storage as storage
+else:
+    import zarr
+    from zarr import storage

asdf_zarr/converter.py

@@ -2,8 +2,7 @@
 
 import asdf
 
-import zarr
-
+from ._zarr_compat import zarr
 from . import util
 from . import storage
 
@@ -12,7 +11,10 @@
 
 class ZarrConverter(asdf.extension.Converter):
     tags = ["asdf://stsci.edu/example-project/tags/zarr-*"]
-    types = ["zarr.core.Array"]
+    types = [
+        "zarr.core.Array",
+        "zarr.v2.core.Array"
+    ]
 
     def to_yaml_tree(self, obj, tag, ctx):
         chunk_store = obj.chunk_store or obj.store
@@ -52,11 +54,6 @@ def to_yaml_tree(self, obj, tag, ctx):
         return obj_dict
 
     def from_yaml_tree(self, node, tag, ctx):
-        import zarr
-
-        from . import util
-        from . import storage
-
         if ".zarray" in node and "chunk_block_map" in node:
             # this is an internally stored zarr array
             # TODO should we enforce no zarr compression here?

asdf_zarr/storage.py

@@ -4,7 +4,8 @@
 
 import asdf
 import numpy
-import zarr
+
+from ._zarr_compat import zarr
 
 
 MISSING_CHUNK = -1

asdf_zarr/tests/test_zarr.py

@@ -6,8 +6,9 @@
 import asdf_zarr.storage
 import numpy
 import pytest
-import zarr
-from zarr.storage import DirectoryStore, KVStore, MemoryStore, NestedDirectoryStore, TempStore
+
+from asdf_zarr._zarr_compat import storage
+from asdf_zarr._zarr_compat import zarr
 
 
 def create_zarray(shape=None, chunks=None, dtype="f8", store=None, chunk_store=None):
@@ -29,14 +30,14 @@ def create_zarray(shape=None, chunks=None, dtype="f8", store=None, chunk_store=N
 @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, KVStore, MemoryStore, NestedDirectoryStore, TempStore])
+@pytest.mark.parametrize("store_type", [storage.DirectoryStore, storage.KVStore, storage.MemoryStore, storage.NestedDirectoryStore, storage.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):
+    if store_type in (storage.DirectoryStore, storage.NestedDirectoryStore):
         store1 = store_type(tmp_path / "zarr_array_1")
         store2 = store_type(tmp_path / "zarr_array_2")
-    elif store_type is KVStore:
+    elif store_type is storage.KVStore:
         store1 = store_type({})
         store2 = store_type({})
     else:
@@ -46,9 +47,9 @@ def test_write_to(tmp_path, copy_arrays, lazy_load, compression, store_type, to_
     # should meta be in a different store?
     if meta_store:
         chunk_store1 = store1
-        store1 = KVStore({})
+        store1 = storage.KVStore({})
         chunk_store2 = store2
-        store2 = KVStore({})
+        store2 = storage.KVStore({})
     else:
         chunk_store1 = None
         chunk_store2 = None
@@ -69,7 +70,7 @@ def test_write_to(tmp_path, copy_arrays, lazy_load, compression, store_type, to_
     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)
-            if to_internal or store_type in (KVStore, MemoryStore, TempStore):
+            if to_internal or store_type in (storage.KVStore, storage.MemoryStore, storage.TempStore):
                 assert isinstance(af[n].chunk_store, asdf_zarr.storage.InternalStore)
             else:
                 assert isinstance(af[n].chunk_store, store_type)
@@ -81,7 +82,7 @@ def test_write_to(tmp_path, copy_arrays, lazy_load, compression, store_type, to_
 @pytest.mark.parametrize("with_update", [True, False])
 def test_modify(tmp_path, with_update, copy_arrays, lazy_load):
     # make a file
-    store = DirectoryStore(tmp_path / "zarr_array")
+    store = storage.DirectoryStore(tmp_path / "zarr_array")
     arr = create_zarray(store=store)
     tree = {"arr": arr}
     fn = tmp_path / "test.asdf"
@@ -112,7 +113,7 @@ class CustomStore(zarr.storage.Store, UserDict):
 @pytest.mark.skip("ASDF Converters aren't aware of the open mode")
 @pytest.mark.parametrize("mode", ["r", "rw"])
 def test_open_mode(tmp_path, mode):
-    store = DirectoryStore(tmp_path / "zarr_array")
+    store = storage.DirectoryStore(tmp_path / "zarr_array")
     arr = create_zarray(store=store)
     tree = {"arr": arr}
     fn = tmp_path / "test.asdf"
@@ -134,14 +135,14 @@ def test_open_mode(tmp_path, mode):
 @pytest.mark.parametrize("meta_store", [True, False])
 def test_to_internal(meta_store):
     if meta_store:
-        zarr = create_zarray(store=KVStore({}), chunk_store=TempStore())
+        zarr = create_zarray(store=storage.KVStore({}), chunk_store=storage.TempStore())
     else:
-        zarr = create_zarray(store=TempStore())
+        zarr = create_zarray(store=storage.TempStore())
     internal = asdf_zarr.storage.to_internal(zarr)
     assert isinstance(internal.chunk_store, asdf_zarr.storage.InternalStore)
     # the store shouldn't be wrapped if it's not used for chunks
     if zarr.store is not zarr.chunk_store:
-        assert isinstance(internal.store, KVStore)
+        assert isinstance(internal.store, storage.KVStore)
     # calling it a second time shouldn't re-wrap the store
     same = asdf_zarr.storage.to_internal(internal)
     assert same.chunk_store is internal.chunk_store

asdf_zarr/util.py

@@ -1,8 +1,9 @@
 import copy
 
 import fsspec
-import zarr.storage
-from zarr.storage import DirectoryStore, FSStore, KVStore, NestedDirectoryStore, TempStore
+
+from ._zarr_compat import zarr
+from ._zarr_compat import storage
 
 
 def encode_storage(store):
@@ -20,20 +21,20 @@ def encode_storage(store):
     obj_dict : dictionary encoding
     """
     obj_dict = {"type_string": store.__class__.__name__}
-    if isinstance(store, (DirectoryStore, NestedDirectoryStore)) and not isinstance(store, TempStore):
+    if isinstance(store, (storage.DirectoryStore, storage.NestedDirectoryStore)) and not isinstance(store, storage.TempStore):
         # dimension separator is _dimension separator and should be
         # read from the zarray itself, not the store
         obj_dict["normalize_keys"] = store.normalize_keys
         obj_dict["path"] = store.path
-    elif isinstance(store, FSStore):
+    elif isinstance(store, storage.FSStore):
         obj_dict["normalize_keys"] = store.normalize_keys
         # store.path path within the filesystem
         obj_dict["path"] = store.path
         # store.mode access mode
         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):
+    elif isinstance(store, storage.KVStore):
         obj_dict["map"] = {k: store[k] for k in store}
     else:
         raise NotImplementedError(f"zarr.storage.Store subclass {store.__class__} not supported")
@@ -56,11 +57,11 @@ def decode_storage(obj_dict):  # TODO needs kwargs for dimension sep?
     kwargs = copy.deepcopy(obj_dict)
     args = []
     type_string = kwargs.pop("type_string")
-    if not hasattr(zarr.storage, type_string):
+    if not hasattr(storage, type_string):
         raise NotImplementedError(f"zarr.storage.Store subclass {type_string} not supported")
     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)
+    return getattr(storage, type_string)(*args, **kwargs)