Skip to content

Commit 0b54940

Browse files
thewtexthewtex
committed
WIP: ENH: Add RFC 3 support
1 parent 6271079 commit 0b54940

15 files changed

+324
-106
lines changed

‎ngff_zarr/_zarr_kwargs.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
import zarr
2+
from packaging import version
3+
4+
zarr_version = version.parse(zarr.__version__)
5+
if zarr_version >= version.parse("3.0.0b1"):
6+
zarr_kwargs = {}
7+
else:
8+
zarr_kwargs = {"dimension_separator": "/"}

‎ngff_zarr/cli.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
from .to_ngff_image import to_ngff_image
4343
from .to_ngff_zarr import to_ngff_zarr
4444
from .zarr_metadata import is_unit_supported
45+
from ._zarr_kwargs import zarr_kwargs
4546

4647

4748
def _multiscales_to_ngff_zarr(
@@ -235,9 +236,7 @@ def main():
235236
cache_dir = Path(args.cache_dir).resolve()
236237
if not cache_dir.exists():
237238
Path.makedirs(cache_dir, parents=True)
238-
config.cache_store = zarr.storage.DirectoryStore(
239-
cache_dir, dimension_separator="/"
240-
)
239+
config.cache_store = zarr.storage.DirectoryStore(cache_dir, **zarr_kwargs)
241240

242241
console = Console()
243242
progress = RichProgress(
@@ -304,7 +303,7 @@ def shutdown_client(sig_id, frame): # noqa: ARG001
304303
)
305304
output_store = None
306305
if args.output and output_backend is ConversionBackend.NGFF_ZARR:
307-
output_store = DirectoryStore(args.output, dimension_separator="/")
306+
output_store = DirectoryStore(args.output, **zarr_kwargs)
308307

309308
subtitle = "[red]generation"
310309
if not args.output:

‎ngff_zarr/config.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,9 @@
22
from pathlib import Path
33

44
import dask.config
5-
import zarr
65
from platformdirs import user_cache_dir
76
from zarr.storage import StoreLike
7+
from ._zarr_kwargs import zarr_kwargs
88

99
if dask.config.get("temporary-directory") is not None:
1010
_store_dir = dask.config.get("temporary-directory")
@@ -13,7 +13,14 @@
1313

1414

1515
def default_store_factory():
16-
return zarr.storage.DirectoryStore(_store_dir, dimension_separator="/")
16+
try:
17+
from zarr.storage import DirectoryStore
18+
19+
return DirectoryStore(_store_dir, **zarr_kwargs)
20+
except ImportError:
21+
from zarr.storage import LocalStore
22+
23+
return LocalStore(_store_dir)
1724

1825

1926
try:

‎ngff_zarr/from_ngff_zarr.py‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,13 @@
44

55
import dask.array
66
import zarr
7-
from zarr.storage import BaseStore
7+
import zarr.storage
8+
9+
# Zarr Python 3
10+
if hasattr(zarr.storage, "StoreLike"):
11+
StoreLike = zarr.storage.StoreLike
12+
else:
13+
StoreLike = Union[MutableMapping, str, Path, zarr.storage.BaseStore]
814

915
from .ngff_image import NgffImage
1016
from .to_multiscales import Multiscales
@@ -13,13 +19,13 @@
1319

1420

1521
def from_ngff_zarr(
16-
store: Union[MutableMapping, str, Path, BaseStore],
22+
store: StoreLike,
1723
validate: bool = False,
1824
) -> Multiscales:
1925
"""
2026
Read an OME-Zarr NGFF Multiscales data structure from a Zarr store.
2127
22-
store : MutableMapping, str or Path, zarr.storage.BaseStore
28+
store : StoreLike
2329
Store or path to directory in file system.
2430
2531
validate : bool

‎ngff_zarr/to_multiscales.py‎

Lines changed: 27 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,12 @@
1111
import zarr
1212
from dask.array.core import Array as DaskArray
1313
from numpy.typing import ArrayLike
14-
from zarr.core import Array as ZarrArray
14+
15+
try:
16+
from zarr.core import Array as ZarrArray
17+
except ImportError:
18+
from zarr.core.array import Array as ZarrArray
19+
from ._zarr_kwargs import zarr_kwargs
1520

1621
from .config import config
1722
from .memory_usage import memory_usage
@@ -82,12 +87,22 @@ def _large_image_serialization(
8287
def remove_from_cache_store(sig_id, frame): # noqa: ARG001
8388
nonlocal base_path_removed
8489
if not base_path_removed:
85-
if isinstance(cache_store, zarr.storage.DirectoryStore):
86-
full_path = Path(cache_store.dir_path()) / base_path
87-
if full_path.exists():
88-
shutil.rmtree(full_path, ignore_errors=True)
89-
else:
90-
zarr.storage.rmdir(cache_store, base_path)
90+
try:
91+
from zarr.storage import DirectoryStore
92+
93+
if isinstance(cache_store, zarr.storage.DirectoryStore):
94+
full_path = Path(cache_store.dir_path()) / base_path
95+
if full_path.exists():
96+
shutil.rmtree(full_path, ignore_errors=True)
97+
else:
98+
zarr.storage.rmdir(cache_store, base_path)
99+
except ImportError:
100+
if isinstance(cache_store, zarr.storage.LocalStore):
101+
full_path = Path(cache_store.root) / base_path
102+
if full_path.exists():
103+
shutil.rmtree(full_path, ignore_errors=True)
104+
else:
105+
zarr.storage.rmdir(cache_store, base_path)
91106
base_path_removed = True
92107

93108
atexit.register(remove_from_cache_store, None, None)
@@ -136,7 +151,7 @@ def remove_from_cache_store(sig_id, frame): # noqa: ARG001
136151
store=cache_store,
137152
path=path,
138153
mode="a",
139-
dimension_separator="/",
154+
**zarr_kwargs,
140155
)
141156

142157
n_slabs = int(np.ceil(data.shape[z_index] / slab_slices))
@@ -164,7 +179,7 @@ def remove_from_cache_store(sig_id, frame): # noqa: ARG001
164179
overwrite=False,
165180
compute=True,
166181
return_stored=False,
167-
dimension_separator="/",
182+
**zarr_kwargs,
168183
)
169184
data = dask.array.from_zarr(cache_store, component=path)
170185
if optimized_chunks < data.shape[z_index] and slab_slices < optimized_chunks:
@@ -180,7 +195,7 @@ def remove_from_cache_store(sig_id, frame): # noqa: ARG001
180195
store=cache_store,
181196
path=path,
182197
mode="a",
183-
dimension_separator="/",
198+
**zarr_kwargs,
184199
)
185200
n_slabs = int(np.ceil(data.shape[z_index] / optimized_chunks))
186201
for slab_index in range(n_slabs):
@@ -205,7 +220,7 @@ def remove_from_cache_store(sig_id, frame): # noqa: ARG001
205220
overwrite=False,
206221
compute=True,
207222
return_stored=False,
208-
dimension_separator="/",
223+
**zarr_kwargs,
209224
)
210225
data = dask.array.from_zarr(cache_store, component=path)
211226
else:
@@ -223,7 +238,7 @@ def remove_from_cache_store(sig_id, frame): # noqa: ARG001
223238
overwrite=False,
224239
compute=True,
225240
return_stored=False,
226-
dimension_separator="/",
241+
**zarr_kwargs,
227242
)
228243
data = dask.array.from_zarr(cache_store, component=path)
229244

‎ngff_zarr/to_ngff_image.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,11 @@
44
import dask
55
from dask.array.core import Array as DaskArray
66
from numpy.typing import ArrayLike
7-
from zarr.core import Array as ZarrArray
7+
8+
try:
9+
from zarr.core import Array as ZarrArray
10+
except ImportError:
11+
from zarr.core.array import Array as ZarrArray
812

913
from .methods._support import _spatial_dims
1014
from .ngff_image import NgffImage

‎ngff_zarr/to_ngff_zarr.py‎

Lines changed: 44 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,18 @@
1111

1212
import dask.array
1313
import numpy as np
14-
import zarr
1514
from itkwasm import array_like_to_numpy_array
16-
from zarr.storage import BaseStore
15+
16+
import zarr
17+
import zarr.storage
18+
19+
# Zarr Python 3
20+
if hasattr(zarr.storage, "StoreLike"):
21+
StoreLike = zarr.storage.StoreLike
22+
else:
23+
StoreLike = Union[MutableMapping, str, Path, zarr.storage.BaseStore]
24+
from ._zarr_kwargs import zarr_kwargs
25+
1726

1827
from .config import config
1928
from .memory_usage import memory_usage
@@ -34,9 +43,7 @@ def _pop_metadata_optionals(metadata_dict):
3443
return metadata_dict
3544

3645

37-
def _prep_for_to_zarr(
38-
store: Union[MutableMapping, str, Path, BaseStore], arr: dask.array.Array
39-
) -> dask.array.Array:
46+
def _prep_for_to_zarr(store: StoreLike, arr: dask.array.Array) -> dask.array.Array:
4047
try:
4148
importlib_metadata.distribution("kvikio")
4249
_KVIKIO_AVAILABLE = True
@@ -81,23 +88,27 @@ def _write_with_tensorstore(store_path: str, array, region, chunks) -> None:
8188

8289

8390
def to_ngff_zarr(
84-
store: Union[MutableMapping, str, Path, BaseStore],
91+
store: StoreLike,
8592
multiscales: Multiscales,
93+
version: str = "0.4",
8694
overwrite: bool = True,
8795
use_tensorstore: bool = False,
88-
chunk_store: Optional[Union[MutableMapping, str, Path, BaseStore]] = None,
96+
chunk_store: Optional[StoreLike] = None,
8997
progress: Optional[Union[NgffProgress, NgffProgressCallback]] = None,
9098
**kwargs,
9199
) -> None:
92100
"""
93101
Write an image pixel array and metadata to a Zarr store with the OME-NGFF standard data model.
94102
95103
:param store: Store or path to directory in file system.
96-
:type store: MutableMapping, str or Path, zarr.storage.BaseStore
104+
:type store: StoreLike
97105
98106
:param multiscales: Multiscales OME-NGFF image pixel data and metadata. Can be generated with ngff_zarr.to_multiscales.
99107
:type multiscales: Multiscales
100108
109+
:param version: OME-Zarr specification version.
110+
:type version: str, optional
111+
101112
:param overwrite: If True, delete any pre-existing data in `store` before creating groups.
102113
:type overwrite: bool, optional
103114
@@ -106,7 +117,7 @@ def to_ngff_zarr(
106117
107118
:param chunk_store: Separate storage for chunks. If not provided, `store` will be used
108119
for storage of both chunks and metadata.
109-
:type chunk_store: MutableMapping, str or Path, zarr.storage.BaseStore, optional
120+
:type chunk_store: StoreLike, optional
110121
111122
:type progress: RichDaskProgress
112123
:param progress: Optional progress logger
@@ -120,10 +131,29 @@ def to_ngff_zarr(
120131
else:
121132
raise ValueError("Tensorstore requires a path-like store")
122133

134+
if version != "0.4" and version != "0.5":
135+
raise ValueError(f"Unsupported version: {version}")
136+
137+
if version == "0.5" and not use_tensorstore:
138+
raise ValueError("Version 0.5 currently requires use_tensorstore=True")
139+
123140
metadata_dict = asdict(multiscales.metadata)
124141
metadata_dict = _pop_metadata_optionals(metadata_dict)
125142
metadata_dict["@type"] = "ngff:Image"
126-
root = zarr.group(store, overwrite=overwrite, chunk_store=chunk_store)
143+
if version == "0.4":
144+
# root = zarr.group(store, overwrite=overwrite, chunk_store=chunk_store)
145+
root = zarr.open_group(
146+
store, mode="w" if overwrite else "a", chunk_store=chunk_store
147+
)
148+
else:
149+
# For version >= 0.5, open root with Zarr v3
150+
root = zarr.open_group(
151+
store,
152+
mode="w" if overwrite else "a",
153+
chunk_store=chunk_store,
154+
zarr_version=3,
155+
)
156+
127157
root.attrs["multiscales"] = [metadata_dict]
128158

129159
nscales = len(multiscales.images)
@@ -160,14 +190,15 @@ def to_ngff_zarr(
160190
shrink_factors.append(1)
161191

162192
chunks = tuple([c[0] for c in arr.chunks])
193+
163194
zarr_array = zarr.creation.open_array(
164195
shape=arr.shape,
165196
chunks=chunks,
166197
dtype=arr.dtype,
167198
store=store,
168199
path=path,
169200
mode="a",
170-
dimension_separator="/",
201+
**zarr_kwargs,
171202
)
172203

173204
shape = image.data.shape
@@ -311,7 +342,7 @@ def to_ngff_zarr(
311342
overwrite=False,
312343
compute=True,
313344
return_stored=False,
314-
dimension_separator="/",
345+
**zarr_kwargs,
315346
**kwargs,
316347
)
317348
else:
@@ -334,7 +365,7 @@ def to_ngff_zarr(
334365
overwrite=False,
335366
compute=True,
336367
return_stored=False,
337-
dimension_separator="/",
368+
**zarr_kwargs,
338369
**kwargs,
339370
)
340371

0 commit comments

Comments
 (0)