diff --git a/changes/2856.feature.rst b/changes/2856.feature.rst new file mode 100644 index 0000000000..3b670c589a --- /dev/null +++ b/changes/2856.feature.rst @@ -0,0 +1 @@ +Opening a string or Path path is redirected to ZipStore if the path has a .zip suffix. \ No newline at end of file diff --git a/src/zarr/storage/_common.py b/src/zarr/storage/_common.py index d81369f142..6fa03cd8d7 100644 --- a/src/zarr/storage/_common.py +++ b/src/zarr/storage/_common.py @@ -11,6 +11,7 @@ from zarr.storage._local import LocalStore from zarr.storage._memory import MemoryStore from zarr.storage._utils import normalize_path +from zarr.storage._zip import ZipStore if TYPE_CHECKING: from zarr.core.buffer import BufferPrototype @@ -241,7 +242,8 @@ async def make_store_path( `StoreLike` object can be a `Store`, `StorePath`, `Path`, `str`, or `dict[str, Buffer]`. If the `StoreLike` object is a Store or `StorePath`, it is converted to a `StorePath` object. If the `StoreLike` object is a Path or str, it is converted - to a LocalStore object and then to a `StorePath` object. If the `StoreLike` + to a LocalStore object and then to a `StorePath` object, unless it has a .zip suffix, + in which case a ZipStore object is used to create the `StorePath`. If the `StoreLike` object is a dict[str, Buffer], it is converted to a `MemoryStore` object and then to a `StorePath` object. @@ -290,11 +292,15 @@ async def make_store_path( else: assert mode in (None, "r", "r+", "a", "w", "w-") # if mode 'r' was provided, we'll open any new stores as read-only + if mode is None: + mode = "r" _read_only = mode == "r" if isinstance(store_like, Store): store = store_like elif store_like is None: store = await MemoryStore.open(read_only=_read_only) + elif isinstance(store_like, Path) and store_like.suffix == ".zip": + store = await ZipStore.open(path=store_like, mode=mode) elif isinstance(store_like, Path): store = await LocalStore.open(root=store_like, read_only=_read_only) elif isinstance(store_like, str): @@ -305,6 +311,8 @@ async def make_store_path( store = FsspecStore.from_url( store_like, storage_options=storage_options, read_only=_read_only ) + elif store_like.endswith(".zip"): + store = await ZipStore.open(path=Path(store_like), mode=mode) else: store = await LocalStore.open(root=Path(store_like), read_only=_read_only) elif isinstance(store_like, dict): diff --git a/tests/test_store/test_core.py b/tests/test_store/test_core.py index bce582a746..349a4d9ef8 100644 --- a/tests/test_store/test_core.py +++ b/tests/test_store/test_core.py @@ -1,12 +1,13 @@ import tempfile from pathlib import Path +import numpy as np import pytest from _pytest.compat import LEGACY_PATH -from zarr import Group +from zarr import Group, open_group from zarr.core.common import AccessModeLiteral, ZarrFormat -from zarr.storage import FsspecStore, LocalStore, MemoryStore, StoreLike, StorePath +from zarr.storage import FsspecStore, LocalStore, MemoryStore, StoreLike, StorePath, ZipStore from zarr.storage._common import contains_array, contains_group, make_store_path from zarr.storage._utils import _join_paths, _normalize_path_keys, _normalize_paths, normalize_path @@ -83,6 +84,36 @@ async def test_make_store_path_local( assert store_path.read_only == (mode == "r") +@pytest.mark.parametrize("store_type", [str, Path]) +@pytest.mark.parametrize("mode", ["r", "w"]) +async def test_make_store_path_zip_path( + tmpdir: LEGACY_PATH, + store_type: type[str] | type[Path] | type[LocalStore], + mode: AccessModeLiteral, +) -> None: + """ + Test that make_store_path creates a ZipStore given a path ending in .zip + """ + zippath = Path(tmpdir) / "zarr.zip" + store_like = store_type(str(zippath)) + + if mode == "r": + store = ZipStore(zippath, mode="w") + root = open_group(store=store, mode="w") + data = np.arange(10000, dtype=np.uint16).reshape(100, 100) + z = root.create_array( + shape=data.shape, chunks=(10, 10), name="foo", dtype=np.uint16, fill_value=99 + ) + z[:] = data + store.close() + + store_path = await make_store_path(store_like, mode=mode) + assert isinstance(store_path.store, ZipStore) + assert Path(store_path.store.path) == zippath + assert store_path.path == normalize_path("") + assert store_path.read_only == (mode == "r") + + @pytest.mark.parametrize("path", [None, "", "bar"]) @pytest.mark.parametrize("mode", ["r", "w"]) async def test_make_store_path_store_path(