Clean up public store API (#2603)

* Clean up public store API

* fix import

Author: dstansby

src/zarr/core/array.py

@@ -113,7 +113,7 @@
     get_pipeline_class,
 )
 from zarr.storage import StoreLike, make_store_path
-from zarr.storage.common import StorePath, ensure_no_existing_node
+from zarr.storage._common import StorePath, ensure_no_existing_node
 
 if TYPE_CHECKING:
     from collections.abc import Iterator, Sequence

src/zarr/core/group.py

@@ -50,8 +50,8 @@
 from zarr.core.metadata.v3 import V3JsonEncoder
 from zarr.core.sync import SyncMixin, sync
 from zarr.errors import MetadataValidationError
-from zarr.storage import StoreLike, make_store_path
-from zarr.storage.common import StorePath, ensure_no_existing_node
+from zarr.storage import StoreLike, StorePath, make_store_path
+from zarr.storage._common import ensure_no_existing_node
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Generator, Iterable, Iterator

src/zarr/storage/__init__.py

@@ -3,16 +3,17 @@
 from types import ModuleType
 from typing import Any
 
-from zarr.storage.common import StoreLike, StorePath, make_store_path
-from zarr.storage.fsspec import FsspecStore
-from zarr.storage.local import LocalStore
-from zarr.storage.logging import LoggingStore
-from zarr.storage.memory import MemoryStore
-from zarr.storage.wrapper import WrapperStore
-from zarr.storage.zip import ZipStore
+from zarr.storage._common import StoreLike, StorePath, make_store_path
+from zarr.storage._fsspec import FsspecStore
+from zarr.storage._local import LocalStore
+from zarr.storage._logging import LoggingStore
+from zarr.storage._memory import GpuMemoryStore, MemoryStore
+from zarr.storage._wrapper import WrapperStore
+from zarr.storage._zip import ZipStore
 
 __all__ = [
     "FsspecStore",
+    "GpuMemoryStore",
     "LocalStore",
     "LoggingStore",
     "MemoryStore",

src/zarr/storage/common.py renamed to src/zarr/storage/_common.py

@@ -8,9 +8,9 @@
 from zarr.core.buffer import Buffer, default_buffer_prototype
 from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZGROUP_JSON, AccessModeLiteral, ZarrFormat
 from zarr.errors import ContainsArrayAndGroupError, ContainsArrayError, ContainsGroupError
+from zarr.storage._local import LocalStore
+from zarr.storage._memory import MemoryStore
 from zarr.storage._utils import normalize_path
-from zarr.storage.local import LocalStore
-from zarr.storage.memory import MemoryStore
 
 if TYPE_CHECKING:
     from zarr.core.buffer import BufferPrototype
@@ -281,7 +281,7 @@ async def make_store_path(
     TypeError
         If the StoreLike object is not one of the supported types.
     """
-    from zarr.storage.fsspec import FsspecStore  # circular import
+    from zarr.storage._fsspec import FsspecStore  # circular import
 
     used_storage_options = False
     path_normalized = normalize_path(path)

src/zarr/storage/fsspec.py renamed to src/zarr/storage/_fsspec.py

@@ -4,7 +4,7 @@
 from typing import TYPE_CHECKING, Any
 
 from zarr.abc.store import ByteRangeRequest, Store
-from zarr.storage.common import _dereference_path
+from zarr.storage._common import _dereference_path
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Iterable

src/zarr/storage/local.py renamed to src/zarr/storage/_local.py

@@ -8,7 +8,7 @@
 from typing import TYPE_CHECKING, Any
 
 from zarr.abc.store import Store
-from zarr.storage.wrapper import WrapperStore
+from zarr.storage._wrapper import WrapperStore
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Generator, Iterable

src/zarr/storage/logging.py renamed to src/zarr/storage/_logging.py

@@ -4,7 +4,7 @@
 import pickle
 from typing import TYPE_CHECKING, Generic, TypeVar
 
-from zarr.storage.wrapper import WrapperStore
+from zarr.storage import WrapperStore
 
 if TYPE_CHECKING:
     from typing import Any

src/zarr/storage/memory.py renamed to src/zarr/storage/_memory.py

@@ -11,7 +11,7 @@
 from zarr.core.common import ZarrFormat
 from zarr.core.sync import sync
 from zarr.storage import MemoryStore, StoreLike
-from zarr.storage.common import _dereference_path
+from zarr.storage._common import _dereference_path
 
 # Copied from Xarray
 _attr_keys = st.text(st.characters(), min_size=1)

src/zarr/storage/wrapper.py renamed to src/zarr/storage/_wrapper.py

@@ -12,8 +12,7 @@
 from zarr import AsyncGroup, config
 from zarr.abc.store import Store
 from zarr.core.sync import sync
-from zarr.storage import LocalStore, MemoryStore, StorePath, ZipStore
-from zarr.storage.fsspec import FsspecStore
+from zarr.storage import FsspecStore, LocalStore, MemoryStore, StorePath, ZipStore
 
 if TYPE_CHECKING:
     from collections.abc import Generator

src/zarr/storage/zip.py renamed to src/zarr/storage/_zip.py

@@ -25,8 +25,8 @@
 )
 from zarr.core.common import JSON, MemoryOrder, ZarrFormat
 from zarr.errors import MetadataValidationError
+from zarr.storage import MemoryStore
 from zarr.storage._utils import normalize_path
-from zarr.storage.memory import MemoryStore
 
 
 def test_create(memory_store: Store) -> None:

src/zarr/testing/store.py

@@ -40,8 +40,7 @@
 from zarr.core.metadata.v3 import DataType
 from zarr.core.sync import sync
 from zarr.errors import ContainsArrayError, ContainsGroupError
-from zarr.storage import LocalStore, MemoryStore
-from zarr.storage.common import StorePath
+from zarr.storage import LocalStore, MemoryStore, StorePath
 
 if TYPE_CHECKING:
     from zarr.core.array_spec import ArrayConfigLike

src/zarr/testing/strategies.py

@@ -12,8 +12,7 @@
 from zarr.codecs.transpose import TransposeCodec
 from zarr.codecs.zstd import ZstdCodec
 from zarr.core.buffer import ArrayLike, BufferPrototype, NDArrayLike, cpu, gpu
-from zarr.storage.common import StorePath
-from zarr.storage.memory import MemoryStore
+from zarr.storage import MemoryStore, StorePath
 from zarr.testing.buffer import (
     NDBufferUsingTestNDArrayLike,
     StoreExpectingTestBuffer,

tests/conftest.py

@@ -7,7 +7,7 @@
 from zarr.abc.store import Store
 from zarr.codecs import BloscCodec
 from zarr.core.buffer import default_buffer_prototype
-from zarr.storage.common import StorePath
+from zarr.storage import StorePath
 
 
 @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])

tests/test_api.py

@@ -6,7 +6,7 @@
 import zarr
 from zarr.abc.store import Store
 from zarr.codecs import BytesCodec
-from zarr.storage.common import StorePath
+from zarr.storage import StorePath
 
 from .test_codecs import _AsyncArrayProxy
 

tests/test_array.py

@@ -4,7 +4,7 @@
 import zarr
 from zarr.abc.store import Store
 from zarr.codecs import GzipCodec
-from zarr.storage.common import StorePath
+from zarr.storage import StorePath
 
 
 @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])

tests/test_buffer.py

@@ -17,7 +17,7 @@
     TransposeCodec,
 )
 from zarr.core.buffer import default_buffer_prototype
-from zarr.storage.common import StorePath
+from zarr.storage import StorePath
 
 from ..conftest import ArrayRequest
 from .test_codecs import _AsyncArrayProxy, order_from_dim

tests/test_codecs/test_blosc.py

@@ -6,7 +6,7 @@
 from zarr.abc.store import Store
 from zarr.codecs import TransposeCodec
 from zarr.core.common import MemoryOrder
-from zarr.storage.common import StorePath
+from zarr.storage import StorePath
 
 from .test_codecs import _AsyncArrayProxy
 

tests/test_codecs/test_endian.py

@@ -10,7 +10,7 @@
 from zarr.codecs import ZstdCodec
 from zarr.core.metadata.v3 import ArrayV3Metadata, DataType
 from zarr.core.strings import _NUMPY_SUPPORTS_VLEN_STRING
-from zarr.storage.common import StorePath
+from zarr.storage import StorePath
 
 numpy_str_dtypes: list[type | str | None] = [None, str, "str", np.dtypes.StrDType]
 expected_zarr_string_dtype: np.dtype[Any]

tests/test_codecs/test_gzip.py

@@ -4,7 +4,7 @@
 import zarr
 from zarr.abc.store import Store
 from zarr.codecs import ZstdCodec
-from zarr.storage.common import StorePath
+from zarr.storage import StorePath
 
 
 @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])

tests/test_codecs/test_sharding.py

@@ -21,8 +21,7 @@
 from zarr.core.group import ConsolidatedMetadata, GroupMetadata
 from zarr.core.sync import sync
 from zarr.errors import ContainsArrayError, ContainsGroupError
-from zarr.storage import LocalStore, MemoryStore, StorePath, ZipStore
-from zarr.storage.common import make_store_path
+from zarr.storage import LocalStore, MemoryStore, StorePath, ZipStore, make_store_path
 
 from .conftest import parse_store
 

tests/test_codecs/test_transpose.py

@@ -26,8 +26,7 @@
     replace_ellipsis,
 )
 from zarr.registry import get_ndbuffer_class
-from zarr.storage.common import StorePath
-from zarr.storage.memory import MemoryStore
+from zarr.storage import MemoryStore, StorePath
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator

tests/test_codecs/test_vlen.py

@@ -21,7 +21,7 @@
 from zarr.core.group import ConsolidatedMetadata, GroupMetadata
 from zarr.core.metadata import ArrayV3Metadata
 from zarr.core.metadata.v2 import ArrayV2Metadata
-from zarr.storage.common import StorePath
+from zarr.storage import StorePath
 
 if TYPE_CHECKING:
     from zarr.abc.store import Store

tests/test_codecs/test_zstd.py

@@ -5,11 +5,8 @@
 from _pytest.compat import LEGACY_PATH
 
 from zarr.core.common import AccessModeLiteral
+from zarr.storage import FsspecStore, LocalStore, MemoryStore, StoreLike, StorePath, make_store_path
 from zarr.storage._utils import normalize_path
-from zarr.storage.common import StoreLike, StorePath, make_store_path
-from zarr.storage.fsspec import FsspecStore
-from zarr.storage.local import LocalStore
-from zarr.storage.memory import MemoryStore
 
 
 @pytest.mark.parametrize("path", [None, "", "bar"])

tests/test_group.py

@@ -6,7 +6,7 @@
 
 import zarr
 from zarr.core.buffer import Buffer, cpu
-from zarr.storage.local import LocalStore
+from zarr.storage import LocalStore
 from zarr.testing.store import StoreTests
 
 if TYPE_CHECKING:

tests/test_indexing.py

@@ -6,7 +6,7 @@
 
 import zarr
 from zarr.core.buffer import default_buffer_prototype
-from zarr.storage.logging import LoggingStore
+from zarr.storage import LoggingStore
 
 if TYPE_CHECKING:
     from zarr.abc.store import Store

tests/test_metadata/test_consolidated.py

@@ -3,7 +3,7 @@
 import pytest
 
 from zarr.core.buffer import Buffer, cpu, gpu
-from zarr.storage.memory import GpuMemoryStore, MemoryStore
+from zarr.storage import GpuMemoryStore, MemoryStore
 from zarr.testing.store import StoreTests
 from zarr.testing.utils import gpu_test
 

tests/test_store/test_core.py

@@ -5,7 +5,7 @@
 import pytest
 
 from zarr.core.buffer.cpu import Buffer, buffer_prototype
-from zarr.storage.wrapper import WrapperStore
+from zarr.storage import WrapperStore
 
 if TYPE_CHECKING:
     from zarr.abc.store import Store

tests/test_store/test_local.py

@@ -10,7 +10,7 @@
 
 import zarr
 from zarr.core.buffer import Buffer, cpu, default_buffer_prototype
-from zarr.storage.zip import ZipStore
+from zarr.storage import ZipStore
 from zarr.testing.store import StoreTests
 
 if TYPE_CHECKING:

tests/test_store/test_logging.py

@@ -14,7 +14,7 @@
     cleanup_resources,
     sync,
 )
-from zarr.storage.memory import MemoryStore
+from zarr.storage import MemoryStore
 
 
 @pytest.fixture(params=[True, False])