add back support zarr-python v2

pyproject.toml

@@ -39,12 +39,14 @@ dependencies = [
     "requests",
     "aiohttp", 
     "dask[array]",
-    "zarr==v3.0.0-alpha.4",
+    "zarr",
 ]
 
 # https://peps.python.org/pep-0621/#dependencies-optional-dependencies
 # "extras" (e.g. for `pip install .[test]`)
 [project.optional-dependencies]
+v3 = ["zarr==v3.0.0-alpha.4"]
+v2 = ["zarr<3.0.0a0"]
 # add dependencies used for testing here
 test = ["pytest", "pytest-cov"]
 # add anything else you like to have in your dev environment here
@@ -60,6 +62,7 @@ dev = [
     "ruff",
 ]
 
+
 [project.urls]
 homepage = "https://github.com/lorenzocerrone/ngio"
 repository = "https://github.com/lorenzocerrone/ngio"
@@ -168,11 +171,15 @@ platforms = ["osx-arm64"]
 
 [tool.pixi.pypi-dependencies]
 # zarr = { path = "../zarr-python/", editable = true }
+# anndata = { path = "../anndata/", editable = true }
 ngio = { path = ".", editable = true }
 
 [tool.pixi.environments]
 default = { solve-group = "default" }
-dev = { features = ["dev"], solve-group = "default" }
-test = { features = ["test"], solve-group = "default" }
+v3 = { features = ["v3"], solve-group = "v3" }
+v2 = { features = ["v2"], solve-group = "v2" }
+dev2 = { features = ["dev"], solve-group = "v2" }
+dev3 = { features = ["dev"], solve-group = "v3" }
+
+test = { features = ["test"], solve-group = "v2" }
 
-[tool.pixi.tasks]

src/ngio/core/label_handler.py

@@ -1,9 +1,7 @@
 """A module to handle OME-NGFF images stored in Zarr format."""
 
-from zarr.store.common import StoreLike
-
 from ngio.core.image_like_handler import ImageLike
-from ngio.io import StoreOrGroup
+from ngio.io import StoreLike, StoreOrGroup
 from ngio.ngff_meta.fractal_image_meta import LabelMeta, PixelSize
 
 

src/ngio/core/ngff_image.py

@@ -2,10 +2,8 @@
 
 from typing import Protocol, TypeVar
 
-from zarr.store.common import StoreLike
-
 from ngio.core.image_handler import Image
-from ngio.io import open_group_wrapper
+from ngio.io import StoreLike, open_group_wrapper
 from ngio.ngff_meta import FractalImageLabelMeta, get_ngff_image_meta_handler
 
 T = TypeVar("T")

src/ngio/io/__init__.py

@@ -1,13 +1,19 @@
 """Collection of helper functions to work with Zarr groups."""
 
 from zarr import Group
-from zarr.store.common import StoreLike
 
-from ngio.io._zarr_group_utils import StoreOrGroup, open_group_wrapper
+from ngio.io._zarr import AccessModeLiteral, StoreLike, StoreOrGroup
+from ngio.io._zarr_group_utils import (
+    open_group_wrapper,
+)
+
+# Zarr V3 imports
+# from zarr.store.common import StoreLike
 
 __all__ = [
     "Group",
     "StoreLike",
+    "AccessModeLiteral",
     "StoreOrGroup",
     "open_group_wrapper",
 ]

src/ngio/io/_zarr.py

@@ -0,0 +1,71 @@
+from importlib.metadata import version
+from pathlib import Path
+from typing import Literal
+
+import zarr
+from packaging.version import Version
+
+zarr_verison = version("zarr")
+ZARR_PYTHON_V = 2 if Version(zarr_verison) < Version("3.0.0a") else 3
+
+# Zarr v3 Imports
+# import zarr.store
+# from zarr.core.common import AccessModeLiteral, ZarrFormat
+# from zarr.store.common import StoreLike
+
+AccessModeLiteral = Literal["r", "r+", "w", "w-", "a"]
+ZarrFormat = Literal[2, 3]
+StoreLike = str | Path  # This type alias more narrrow than necessary
+StoreOrGroup = StoreLike | zarr.Group
+
+
+class ZarrV3Error(Exception):
+    pass
+
+
+def _pass_through_group(
+    group: zarr.Group, mode: AccessModeLiteral, zarr_format: ZarrFormat = 2
+) -> zarr.Group:
+    if ZARR_PYTHON_V == 2:
+        if zarr_format == 3:
+            raise ZarrV3Error("Zarr v3 is not supported in when using zarr-python v2.")
+        else:
+            return group
+
+    else:
+        if group.metadata.zarr_format != zarr_format:
+            raise ValueError(
+                f"Zarr format mismatch. Expected {zarr_format}, "
+                "got {store.metadata.zarr_format}."
+            )
+        else:
+            return group
+
+    raise ValueError("This should never be reached.")
+
+
+def _open_group_v2_v3(
+    store: StoreOrGroup, mode: AccessModeLiteral, zarr_format: ZarrFormat = 2
+) -> zarr.Group:
+    """Wrapper around zarr.open_group with some additional checks.
+
+    Args:
+        store (StoreOrGroup): The store (can also be a Path/str) or group to open.
+        mode (ReadOrEdirLiteral): The mode to open the group in.
+        zarr_format (ZarrFormat): The Zarr format to use.
+
+    Returns:
+        zarr.Group: The opened Zarr group.
+    """
+    if ZARR_PYTHON_V == 3:
+        return zarr.open_group(store=store, mode=mode, zarr_format=zarr_format)
+    else:
+        return zarr.open_group(store=store, mode=mode, zarr_version=zarr_format)
+
+
+def _is_group_readonly(group: zarr.Group) -> bool:
+    if ZARR_PYTHON_V == 3:
+        return group.store_path.store.mode.readonly
+
+    else:
+        return not group.store.is_writeable()

src/ngio/io/_zarr_group_utils.py

@@ -1,28 +1,29 @@
+from pathlib import Path
+
 import zarr
-import zarr.store
-from zarr.core.common import AccessModeLiteral, ZarrFormat
-from zarr.store.common import StoreLike
 
-StoreOrGroup = StoreLike | zarr.Group
+from ngio.io._zarr import (
+    AccessModeLiteral,
+    StoreLike,
+    StoreOrGroup,
+    ZarrFormat,
+    _open_group_v2_v3,
+    _pass_through_group,
+)
+
+# Zarr v3 Imports
+# import zarr.store
+# from zarr.core.common import AccessModeLiteral, ZarrFormat
+# from zarr.store.common import StoreLike
 
 
 def _check_store(store: StoreLike) -> StoreLike:
-    if isinstance(store, zarr.store.RemoteStore):
-        raise NotImplementedError(
-            "RemoteStore is not yet supported. Please use LocalStore."
-        )
-    return store
+    if isinstance(store, str) or isinstance(store, Path):
+        return store
 
-
-def _pass_through_group(
-    group: zarr.Group, mode: AccessModeLiteral, zarr_format: ZarrFormat = 2
-) -> zarr.Group:
-    if group.metadata.zarr_format != zarr_format:
-        raise ValueError(
-            f"Zarr format mismatch. Expected {zarr_format}, "
-            "got {store.metadata.zarr_format}."
-        )
-    return group
+    raise NotImplementedError(
+        "RemoteStore is not yet supported. Please use LocalStore."
+    )
 
 
 def open_group_wrapper(
@@ -42,4 +43,5 @@ def open_group_wrapper(
         return _pass_through_group(store, mode=mode, zarr_format=zarr_format)
 
     store = _check_store(store)
-    return zarr.open_group(store=store, mode=mode, zarr_format=zarr_format)
+
+    return _open_group_v2_v3(store=store, mode=mode, zarr_format=zarr_format)

src/ngio/ngff_meta/meta_handler.py

@@ -2,9 +2,7 @@
 
 from typing import Literal, Protocol
 
-from zarr.core.common import AccessModeLiteral
-
-from ngio.io import Group, StoreLike, StoreOrGroup
+from ngio.io import AccessModeLiteral, Group, StoreOrGroup
 from ngio.ngff_meta.fractal_image_meta import ImageLabelMeta
 from ngio.ngff_meta.v04.zarr_utils import (
     NgffImageMetaZarrHandlerV04,
@@ -30,7 +28,7 @@ def group(self) -> Group:
         ...
 
     @property
-    def store(self) -> StoreLike:
+    def store(self):
         """Return the Zarr store."""
         ...
 

src/ngio/ngff_meta/v04/zarr_utils.py

@@ -2,14 +2,13 @@
 
 from typing import Literal
 
-from zarr.core.common import AccessModeLiteral
-
 from ngio.io import (
+    AccessModeLiteral,
     Group,
-    StoreLike,
     StoreOrGroup,
     open_group_wrapper,
 )
+from ngio.io._zarr import _is_group_readonly
 from ngio.ngff_meta.fractal_image_meta import (
     Axis,
     Dataset,
@@ -212,7 +211,11 @@ def __init__(
             mode (str): The mode of the store.
         """
         if isinstance(store, Group):
-            self._store = store.store_path
+            if hasattr(store, "store_path"):
+                self._store = store.store_path
+            else:
+                self._store = store.store
+
             self._group = store
 
         else:
@@ -233,7 +236,7 @@ def zarr_version(self) -> int:
         return 2
 
     @property
-    def store(self) -> StoreLike:
+    def store(self):
         """Return the Zarr store."""
         return self._store
 
@@ -263,7 +266,7 @@ def load_meta(self) -> ImageLabelMeta:
 
     def write_meta(self, meta: ImageLabelMeta) -> None:
         """Write the OME-NGFF 0.4 metadata."""
-        if self.group.store_path.store.mode.readonly:
+        if _is_group_readonly(self.group):
             raise ValueError(
                 "The store is read-only. Cannot write the metadata to the store."
             )

tests/core/conftest.py

@@ -1,16 +1,23 @@
 import json
+from importlib.metadata import version
 from pathlib import Path
 
 import zarr
-import zarr.store
+from packaging.version import Version
 from pytest import fixture
 
+zarr_verison = version("zarr")
+ZARR_PYTHON_V = 2 if Version(zarr_verison) < Version("3.0.0a") else 3
+
 
 @fixture
 def ome_zarr_image_v04_path(tmpdir):
     zarr_path = Path(tmpdir) / "test_ome_ngff_v04.zarr"
 
-    group = zarr.open_group(store=zarr_path, mode="w", zarr_format=2)
+    if ZARR_PYTHON_V == 3:
+        group = zarr.open_group(store=zarr_path, mode="w", zarr_format=2)
+    else:
+        group = zarr.open_group(store=zarr_path, mode="w", zarr_version=2)
 
     json_path = (
         Path(".") / "tests" / "data" / "meta_v04" / "base_ome_zarr_image_meta.json"
@@ -24,6 +31,9 @@ def ome_zarr_image_v04_path(tmpdir):
     # shape = (3, 10, 256, 256)
     for i, path in enumerate(["0", "1", "2", "3"]):
         shape = (3, 10, 256 // (2**i), 256 // (2**i))
-        group.create_array(name=path, fill_value=0, shape=shape)
+        if ZARR_PYTHON_V == 3:
+            group.create_array(name=path, fill_value=0, shape=shape)
+        else:
+            group.zeros(name=path, shape=shape)
 
     return zarr_path

tests/io/conftest.py

@@ -1,16 +1,27 @@
+from importlib.metadata import version
 from pathlib import Path
 
 import zarr
-import zarr.store
+from packaging.version import Version
 from pytest import fixture
 
+zarr_verison = version("zarr")
+ZARR_PYTHON_V = 2 if Version(zarr_verison) < Version("3.0.0a") else 3
+
 
 def _create_zarr(tempdir, zarr_format=2):
     zarr_path = Path(tempdir) / f"test_group_v{2}.zarr"
-    group = zarr.open_group(store=zarr_path, mode="w", zarr_format=zarr_format)
+
+    if ZARR_PYTHON_V == 3:
+        group = zarr.open_group(store=zarr_path, mode="w", zarr_format=zarr_format)
+    else:
+        group = zarr.open_group(store=zarr_path, mode="w", zarr_version=zarr_format)
 
     for i in range(3):
-        group.create_array(f"array_{i}", shape=(10, 10), dtype="i4")
+        if ZARR_PYTHON_V == 3:
+            group.create_array(f"array_{i}", shape=(10, 10), dtype="i4")
+        else:
+            group.empty(f"array_{i}", shape=(10, 10), dtype="i4")
 
     for i in range(3):
         group.create_group(f"group_{i}")
@@ -42,19 +53,13 @@ def local_zarr_str_v3(tmpdir) -> tuple[Path, int]:
     return str(zarr_path.absolute()), 3
 
 
-@fixture
-def local_zarr_store_v2(tmpdir) -> zarr.store.LocalStore:
-    zarr_path = _create_zarr(tmpdir, zarr_format=2)
-    return zarr.store.LocalStore(zarr_path, mode="r+"), 2
-
-
 @fixture(
     params=[
         "local_zarr_path_v2",
-        "local_zarr_path_v3",
+        # "local_zarr_path_v3",
         "local_zarr_str_v2",
-        "local_zarr_str_v3",
-        "local_zarr_store_v2",
+        # "local_zarr_str_v3",
+        # "local_zarr_store_v2",
     ]
 )
 def store_fixture(request):

tests/io/test_zarr_group_utils.py

@@ -1,5 +1,6 @@
 import pytest
 import zarr
+from conftest import ZARR_PYTHON_V
 
 
 class TestGroupUtils:
@@ -15,6 +16,7 @@ def test_open_group_wrapper(self, store_fixture):
         group.attrs.update(self.test_attrs)
         assert dict(group.attrs) == self.test_attrs
 
+    @pytest.mark.skipif(ZARR_PYTHON_V, reason="Zarr V2 does not support remote stores.")
     def test_raise_not_implemented_error(self):
         from ngio.io._zarr_group_utils import open_group_wrapper