use Store as input rather than StoreLike

Author: K-Meech

src/zarr/core/metadata/converter/cli.py

@@ -6,6 +6,7 @@
 
 import zarr.core.metadata.converter.migrate_to_v3 as migrate_metadata
 from zarr.core.sync import sync
+from zarr.storage._common import make_store
 
 app = typer.Typer()
 
@@ -100,13 +101,19 @@ def migrate(
             "Dry run enabled - no new files will be created or changed. Log of files that would be created on a real run:"
         )
 
-    write_store = input_store if output_store is None else output_store
+    input_zarr_store = sync(make_store(input_store, mode="r+"))
+
+    if output_store is not None:
+        output_zarr_store = sync(make_store(output_store, mode="w-"))
+        write_store = output_zarr_store
+    else:
+        write_store = input_zarr_store
 
     if overwrite:
         sync(migrate_metadata.remove_metadata(write_store, 3, force=force, dry_run=dry_run))
 
     migrate_metadata.migrate_v2_to_v3(
-        input_store=input_store, output_store=output_store, dry_run=dry_run
+        input_store=input_zarr_store, output_store=output_zarr_store, dry_run=dry_run
     )
 
     if remove_v2_metadata:
@@ -150,10 +157,11 @@ def remove_metadata(
         logger.info(
             "Dry run enabled - no files will be deleted or changed. Log of files that would be deleted on a real run:"
         )
+    input_zarr_store = sync(make_store(store, mode="r+"))
 
     sync(
         migrate_metadata.remove_metadata(
-            store=store,
+            store=input_zarr_store,
             zarr_format=cast(Literal[2, 3], int(zarr_format[1:])),
             force=force,
             dry_run=dry_run,

src/zarr/core/metadata/converter/migrate_to_v3.py

@@ -1,11 +1,12 @@
 import asyncio
 import logging
-from typing import Any, Literal, cast
+from typing import Literal, cast
 
 import numcodecs.abc
 
 import zarr
 from zarr.abc.codec import ArrayArrayCodec, BytesBytesCodec, Codec
+from zarr.abc.store import Store
 from zarr.codecs.blosc import BloscCodec, BloscShuffle
 from zarr.codecs.bytes import BytesCodec
 from zarr.codecs.gzip import GzipCodec
@@ -29,17 +30,15 @@
 from zarr.core.metadata.v3 import ArrayV3Metadata
 from zarr.core.sync import sync
 from zarr.registry import get_codec_class
-from zarr.storage import StoreLike
-from zarr.storage._common import StorePath, make_store_path
+from zarr.storage import StorePath
 
 logger = logging.getLogger(__name__)
 
 
 def migrate_v2_to_v3(
     *,
-    input_store: StoreLike,
-    output_store: StoreLike | None = None,
-    storage_options: dict[str, Any] | None = None,
+    input_store: Store,
+    output_store: Store | None = None,
     dry_run: bool = False,
 ) -> None:
     """Migrate all v2 metadata in a zarr hierarchy to v3.
@@ -49,26 +48,20 @@ def migrate_v2_to_v3(
 
     Parameters
     ----------
-    input_store : StoreLike
-        Input Zarr to migrate - should be a store, path to directory in file system or name of zip file.
-    output_store : StoreLike
+    input_store : Store
+        Input Zarr to migrate.
+    output_store : Store, optional
         Output location to write v3 metadata (no array data will be copied). If not provided, v3 metadata will be
-        written to input_store. Should be a store, path to directory in file system or name of zip file.
-    storage_options : dict | None, optional
-        If the store is backed by an fsspec-based implementation, then this dict will be passed to
-        the Store constructor for that implementation. Ignored otherwise. Note - the same storage_options will
-        be passed to both input_store and output_store (if provided).
+        written to input_store.
     dry_run : bool, optional
         Enable a 'dry run' - files that would be created are logged, but no files are created or changed.
     """
 
-    zarr_v2 = zarr.open(store=input_store, mode="r+", storage_options=storage_options)
+    zarr_v2 = zarr.open(store=input_store, mode="r+")
 
     if output_store is not None:
         # w- access to not allow overwrite of existing data
-        output_path = sync(
-            make_store_path(output_store, mode="w-", storage_options=storage_options)
-        )
+        output_path = sync(StorePath.open(output_store, path="", mode="w-"))
     else:
         output_path = zarr_v2.store_path
 
@@ -101,9 +94,8 @@ def migrate_to_v3(zarr_v2: Array | Group, output_path: StorePath, dry_run: bool
 
 
 async def remove_metadata(
-    store: StoreLike,
+    store: Store,
     zarr_format: ZarrFormat,
-    storage_options: dict[str, Any] | None = None,
     force: bool = False,
     dry_run: bool = False,
 ) -> None:
@@ -113,23 +105,21 @@ async def remove_metadata(
 
     Parameters
     ----------
-    store : StoreLike
-        Store or path to directory in file system or name of zip file.
+    store : Store
+        Zarr to remove metadata from.
     zarr_format : ZarrFormat
         Which format's metadata to remove - 2 or 3.
-    storage_options : dict | None, optional
-        If the store is backed by an fsspec-based implementation, then this dict will be passed to
-        the Store constructor for that implementation. Ignored otherwise.
     force : bool, optional
         When False, metadata can only be removed if a valid alternative exists e.g. deletion of v2 metadata will
         only be allowed when v3 metadata is also present. When True, metadata can be removed when there is no
         alternative.
     dry_run : bool, optional
         Enable a 'dry run' - files that would be deleted are logged, but no files are removed or changed.
     """
-    store_path = await make_store_path(store, mode="r+", storage_options=storage_options)
-    if not store_path.store.supports_deletes:
+
+    if not store.supports_deletes:
         raise ValueError("Store must support deletes to remove metadata")
+    store_path = await StorePath.open(store, path="", mode="r+")
 
     metadata_files_all = {
         2: [ZARRAY_JSON, ZATTRS_JSON, ZGROUP_JSON, ZMETADATA_V2_JSON],
@@ -142,7 +132,7 @@ async def remove_metadata(
         alternative_metadata = 2
 
     awaitables = []
-    async for file_path in store_path.store.list():
+    async for file_path in store.list():
         parent_path, _, file_name = file_path.rpartition("/")
 
         if file_name not in metadata_files_all[zarr_format]:

src/zarr/storage/_common.py

@@ -267,35 +267,105 @@ def __eq__(self, other: object) -> bool:
 StoreLike: TypeAlias = Store | StorePath | FSMap | Path | str | dict[str, Buffer]
 
 
-async def make_store_path(
+async def make_store(
     store_like: StoreLike | None,
     *,
-    path: str | None = "",
     mode: AccessModeLiteral | None = None,
     storage_options: dict[str, Any] | None = None,
-) -> StorePath:
+) -> Store:
     """
-    Convert a `StoreLike` object into a StorePath object.
+    Convert a `StoreLike` object into a Store object.
+
+    `StoreLike` objects are converted to `Store` as follows:
+
+    - `Store` or `StorePath` = `Store` object.
+    - `Path` or `str` = `LocalStore` object.
+    - `str` that starts with a protocol = `FsspecStore` object.
+    - `dict[str, Buffer]` = `MemoryStore` object.
+    - `None` = `MemoryStore` object.
+    - `FSMap` = `FsspecStore` object.
+
+    Parameters
+    ----------
+    store_like : StoreLike | None
+        The object to convert to a `Store` object.
+    mode : StoreAccessMode | None, optional
+        The mode to use when creating the `Store` object.  If None, the
+        default mode is 'r'.
+    storage_options : dict[str, Any] | None, optional
+        The storage options to use when creating the `RemoteStore` object.  If
+        None, the default storage options are used.
+
+    Returns
+    -------
+    Store
+        The converted Store object.
+
+    Raises
+    ------
+    TypeError
+        If the StoreLike object is not one of the supported types, or if storage_options is provided but not used.
+    ValueError
+        If storage_options is provided for a store that does not support it.
+    """
+    from zarr.storage._fsspec import FsspecStore  # circular import
 
-    This function takes a `StoreLike` object and returns a `StorePath` object.  The
-    `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`
-    object is a dict[str, Buffer], it is converted to a `MemoryStore` object and
-    then to a `StorePath` object.
+    used_storage_options = False
+    assert mode in (None, "r", "r+", "a", "w", "w-")
+
+    # if mode 'r' was provided, we'll open any new stores as read-only
+    _read_only = mode == "r"
 
-    If the `StoreLike` object is None, a `MemoryStore` object is created and
-    converted to a `StorePath` object.
+    if isinstance(store_like, StorePath):
+        store = store_like.store
+    elif 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):
+        store = await LocalStore.open(root=store_like, read_only=_read_only)
+    elif isinstance(store_like, str):
+        storage_options = storage_options or {}
+
+        if _is_fsspec_uri(store_like):
+            used_storage_options = True
+            store = FsspecStore.from_url(
+                store_like, storage_options=storage_options, read_only=_read_only
+            )
+        else:
+            store = await LocalStore.open(root=Path(store_like), read_only=_read_only)
+    elif isinstance(store_like, dict):
+        # We deliberate only consider dict[str, Buffer] here, and not arbitrary mutable mappings.
+        # By only allowing dictionaries, which are in-memory, we know that MemoryStore appropriate.
+        store = await MemoryStore.open(store_dict=store_like, read_only=_read_only)
+    elif _has_fsspec and isinstance(store_like, FSMap):
+        if storage_options:
+            raise ValueError(
+                "'storage_options was provided but is not used for FSMap store_like objects. Specify the storage options when creating the FSMap instance instead."
+            )
+        store = FsspecStore.from_mapper(store_like, read_only=_read_only)
+    else:
+        raise TypeError(f"Unsupported type for store_like: '{type(store_like).__name__}'")
 
-    If the `StoreLike` object is a str and starts with a protocol, it is
-    converted to a RemoteStore object and then to a `StorePath` object.
+    if storage_options and not used_storage_options:
+        msg = "'storage_options' was provided but unused. 'storage_options' is only used for fsspec filesystem stores."
+        raise TypeError(msg)
 
-    If the `StoreLike` object is a dict[str, Buffer] and the mode is not None,
-    the `MemoryStore` object is created with the given mode.
+    return store
 
-    If the `StoreLike` object is a str and starts with a protocol, the
-    RemoteStore object is created with the given mode and storage options.
+
+async def make_store_path(
+    store_like: StoreLike | None,
+    *,
+    path: str | None = "",
+    mode: AccessModeLiteral | None = None,
+    storage_options: dict[str, Any] | None = None,
+) -> StorePath:
+    """
+    Convert a `StoreLike` object into a StorePath object.
+
+    This function takes a `StoreLike` object and returns a `StorePath` object. See `make_store` for details
+    of which `Store` is used for each type of `store_like` object.
 
     Parameters
     ----------
@@ -319,58 +389,28 @@ async def make_store_path(
     Raises
     ------
     TypeError
-        If the StoreLike object is not one of the supported types.
-    """
-    from zarr.storage._fsspec import FsspecStore  # circular import
+        If the StoreLike object is not one of the supported types, or if storage_options is provided but not used.
+    ValueError
+        If storage_options is provided for a store that does not support it.
 
-    used_storage_options = False
+    See Also
+    --------
+    make_store
+    """
     path_normalized = normalize_path(path)
+
     if isinstance(store_like, StorePath):
-        result = store_like / path_normalized
+        if storage_options:
+            msg = "'storage_options' was provided but unused. 'storage_options' is only used for fsspec filesystem stores."
+            raise TypeError(msg)
+        return store_like / path_normalized
+    elif _has_fsspec and isinstance(store_like, FSMap) and path:
+        raise ValueError(
+            "'path' was provided but is not used for FSMap store_like objects. Specify the path when creating the FSMap instance instead."
+        )
     else:
-        assert mode in (None, "r", "r+", "a", "w", "w-")
-        # if mode 'r' was provided, we'll open any new stores as read-only
-        _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):
-            store = await LocalStore.open(root=store_like, read_only=_read_only)
-        elif isinstance(store_like, str):
-            storage_options = storage_options or {}
-
-            if _is_fsspec_uri(store_like):
-                used_storage_options = True
-                store = FsspecStore.from_url(
-                    store_like, storage_options=storage_options, read_only=_read_only
-                )
-            else:
-                store = await LocalStore.open(root=Path(store_like), read_only=_read_only)
-        elif isinstance(store_like, dict):
-            # We deliberate only consider dict[str, Buffer] here, and not arbitrary mutable mappings.
-            # By only allowing dictionaries, which are in-memory, we know that MemoryStore appropriate.
-            store = await MemoryStore.open(store_dict=store_like, read_only=_read_only)
-        elif _has_fsspec and isinstance(store_like, FSMap):
-            if path:
-                raise ValueError(
-                    "'path' was provided but is not used for FSMap store_like objects. Specify the path when creating the FSMap instance instead."
-                )
-            if storage_options:
-                raise ValueError(
-                    "'storage_options was provided but is not used for FSMap store_like objects. Specify the storage options when creating the FSMap instance instead."
-                )
-            store = FsspecStore.from_mapper(store_like, read_only=_read_only)
-        else:
-            raise TypeError(f"Unsupported type for store_like: '{type(store_like).__name__}'")
-
-        result = await StorePath.open(store, path=path_normalized, mode=mode)
-
-    if storage_options and not used_storage_options:
-        msg = "'storage_options' was provided but unused. 'storage_options' is only used for fsspec filesystem stores."
-        raise TypeError(msg)
-
-    return result
+        store = await make_store(store_like, mode=mode, storage_options=storage_options)
+        return await StorePath.open(store, path=path_normalized, mode=mode)
 
 
 def _is_fsspec_uri(uri: str) -> bool: