Skip to content

Open zip store from path with zip suffix #2856

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions changes/2856.feature.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Opening a string or Path path is redirected to ZipStore if the path has a .zip suffix.
10 changes: 9 additions & 1 deletion src/zarr/storage/_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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):
Expand All @@ -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):
Expand Down
35 changes: 33 additions & 2 deletions tests/test_store/test_core.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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(
Expand Down