Add suffix detection in make_store_path to open zip stores from path - add test

Author: aulemahal

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.
 
@@ -295,6 +297,8 @@ async def make_store_path(
             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 or "r")
         elif isinstance(store_like, Path):
             store = await LocalStore.open(root=store_like, read_only=_read_only)
         elif isinstance(store_like, str):
@@ -305,6 +309,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 or "r")
             else:
                 store = await LocalStore.open(root=Path(store_like), read_only=_read_only)
         elif isinstance(store_like, dict):

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 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(