From 31f074c65963b9d5bd2b58951ead5486fde3db0c Mon Sep 17 00:00:00 2001 From: lorenzo Date: Mon, 23 Sep 2024 16:41:42 +0200 Subject: [PATCH] add back support zarr-python v2 --- pyproject.toml | 15 ++++-- src/ngio/core/label_handler.py | 4 +- src/ngio/core/ngff_image.py | 4 +- src/ngio/io/__init__.py | 10 +++- src/ngio/io/_zarr.py | 71 ++++++++++++++++++++++++++++ src/ngio/io/_zarr_group_utils.py | 42 ++++++++-------- src/ngio/ngff_meta/meta_handler.py | 6 +-- src/ngio/ngff_meta/v04/zarr_utils.py | 15 +++--- tests/core/conftest.py | 16 +++++-- tests/io/conftest.py | 29 +++++++----- tests/io/test_zarr_group_utils.py | 2 + tests/ngff_meta/conftest.py | 16 +++++-- 12 files changed, 170 insertions(+), 60 deletions(-) create mode 100644 src/ngio/io/_zarr.py diff --git a/pyproject.toml b/pyproject.toml index 43d8aa4..d3b7737 100644 --- a/pyproject.toml +++ b/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] diff --git a/src/ngio/core/label_handler.py b/src/ngio/core/label_handler.py index 1a6d59b..ae85f9e 100644 --- a/src/ngio/core/label_handler.py +++ b/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 diff --git a/src/ngio/core/ngff_image.py b/src/ngio/core/ngff_image.py index c4e8ac2..8f7f61e 100644 --- a/src/ngio/core/ngff_image.py +++ b/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") diff --git a/src/ngio/io/__init__.py b/src/ngio/io/__init__.py index 78649c7..c8e65fa 100644 --- a/src/ngio/io/__init__.py +++ b/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", ] diff --git a/src/ngio/io/_zarr.py b/src/ngio/io/_zarr.py new file mode 100644 index 0000000..d8ef49f --- /dev/null +++ b/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() diff --git a/src/ngio/io/_zarr_group_utils.py b/src/ngio/io/_zarr_group_utils.py index d417fbd..5f37ae7 100644 --- a/src/ngio/io/_zarr_group_utils.py +++ b/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) diff --git a/src/ngio/ngff_meta/meta_handler.py b/src/ngio/ngff_meta/meta_handler.py index 8a8fc93..92f6d43 100644 --- a/src/ngio/ngff_meta/meta_handler.py +++ b/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.""" ... diff --git a/src/ngio/ngff_meta/v04/zarr_utils.py b/src/ngio/ngff_meta/v04/zarr_utils.py index 50f11e9..9ddd95c 100644 --- a/src/ngio/ngff_meta/v04/zarr_utils.py +++ b/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." ) diff --git a/tests/core/conftest.py b/tests/core/conftest.py index 8d66c22..b35de0f 100644 --- a/tests/core/conftest.py +++ b/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 diff --git a/tests/io/conftest.py b/tests/io/conftest.py index aa3d04a..2da1b11 100644 --- a/tests/io/conftest.py +++ b/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): diff --git a/tests/io/test_zarr_group_utils.py b/tests/io/test_zarr_group_utils.py index 433a564..3ef08c9 100644 --- a/tests/io/test_zarr_group_utils.py +++ b/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 diff --git a/tests/ngff_meta/conftest.py b/tests/ngff_meta/conftest.py index f2ed3ba..03a8387 100644 --- a/tests/ngff_meta/conftest.py +++ b/tests/ngff_meta/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) with open("tests/data/meta_v04/base_ome_zarr_image_meta.json") as f: base_ome_zarr_meta = json.load(f) @@ -24,7 +31,10 @@ def ome_zarr_image_v04_path(tmpdir): def ome_zarr_label_v04_path(tmpdir): zarr_path = Path(tmpdir) / "test_ome_ngff_image_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) with open("tests/data/meta_v04/base_ome_zarr_label_meta.json") as f: base_ome_zarr_meta = json.load(f)