add path to stores

Author: d-v-b

src/zarr/abc/store.py

@@ -51,12 +51,17 @@ class Store(ABC):
     #         'attributes')
     #     raise NotImplementedError(msg)
 
-    def __init__(
-        self, path: str = "", mode: AccessModeLiteral = "r", *args: Any, **kwargs: Any
-    ) -> None:
+    def __init__(self, path: str = "", mode: AccessModeLiteral = "r") -> None:
         object.__setattr__(self, "_is_open", False)
         object.__setattr__(self, "_mode", AccessMode.from_literal(mode))
-        object.__setattr__(self, "path", path)
+        object.__setattr__(self, "path", validate_path(path))
+
+    def resolve_key(self, key: str) -> str:
+        key = parse_path(key)
+        if self.path == "":
+            return key
+        else:
+            return f"{self.path}/{key}"
 
     @classmethod
     async def open(cls, *args: Any, **kwargs: Any) -> Self:
@@ -335,8 +340,8 @@ def with_path(self, path: str) -> Self:
         """
         Return a copy of this store with a new path attribute
         """
-        # TODO: implement this
-        return self
+        # TODO: implement me
+        raise NotImplementedError
 
 
 @runtime_checkable
@@ -364,3 +369,20 @@ async def set_or_delete(byte_setter: ByteSetter, value: Buffer | None) -> None:
         await byte_setter.delete()
     else:
         await byte_setter.set(value)
+
+
+def validate_path(path: str) -> str:
+    """
+    Ensure that the input string is a valid relative path in the abstract zarr object storage scheme.
+    """
+    if path.endswith("/"):
+        raise ValueError(f"Invalid path: {path} ends with '/'.")
+    if "//" in path:
+        raise ValueError(f"Invalid path: {path} contains '//'.")
+    if "\\" in path:
+        raise ValueError(f"Invalid path: {path} contains '\"'.")
+    return path
+
+
+def parse_path(path: str) -> str:
+    return path.rstrip("/").lstrip("/")

src/zarr/store/common.py

@@ -101,7 +101,7 @@ async def make_store_path(
             mode = "w"  # exception to the default mode = 'r'
         result = StorePath(await MemoryStore.open(mode=mode))
     elif isinstance(store_like, Path):
-        result = StorePath(await LocalStore.open(root=store_like, mode=mode or "r"))
+        result = StorePath(await LocalStore.open(path=store_like, mode=mode or "r"))
     elif isinstance(store_like, str):
         storage_options = storage_options or {}
 
@@ -111,7 +111,7 @@ async def make_store_path(
                 RemoteStore.from_url(store_like, storage_options=storage_options, mode=mode or "r")
             )
         else:
-            result = StorePath(await LocalStore.open(root=Path(store_like), mode=mode or "r"))
+            result = StorePath(await LocalStore.open(path=store_like, mode=mode or "r"))
     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.

src/zarr/store/local.py

@@ -4,14 +4,15 @@
 import os
 import shutil
 from pathlib import Path
-from typing import TYPE_CHECKING, Self
+from typing import TYPE_CHECKING
 
 from zarr.abc.store import ByteRangeRequest, Store
 from zarr.core.buffer import Buffer
 from zarr.core.common import concurrent_map, to_thread
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Iterable
+    from typing import Self
 
     from zarr.core.buffer import BufferPrototype
     from zarr.core.common import AccessModeLiteral
@@ -125,7 +126,7 @@ async def get(
     ) -> Buffer | None:
         if not self._is_open:
             await self._open()
-        assert isinstance(key, str)
+
         path = os.path.join(self.path, key)
 
         try:

src/zarr/store/logging.py

@@ -5,13 +5,14 @@
 import time
 from collections import defaultdict
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Self
+from typing import TYPE_CHECKING
 
 from zarr.abc.store import AccessMode, ByteRangeRequest, Store
 from zarr.core.buffer import Buffer
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Generator, Iterable
+    from typing import Self
 
     from zarr.core.buffer import Buffer, BufferPrototype
     from zarr.core.common import AccessModeLiteral

src/zarr/store/memory.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Self
+from typing import TYPE_CHECKING
 
 from zarr.abc.store import ByteRangeRequest, Store
 from zarr.core.buffer import Buffer, gpu
@@ -9,6 +9,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Iterable, MutableMapping
+    from typing import Self
 
     from zarr.core.buffer import BufferPrototype
     from zarr.core.common import AccessModeLiteral
@@ -26,9 +27,9 @@ class MemoryStore(Store):
 
     def __init__(
         self,
-        path: str = "",
         store_dict: MutableMapping[str, Buffer] | None = None,
         *,
+        path: str = "",
         mode: AccessModeLiteral = "r",
     ) -> None:
         super().__init__(mode=mode, path=path)
@@ -68,9 +69,9 @@ async def get(
     ) -> Buffer | None:
         if not self._is_open:
             await self._open()
-        assert isinstance(key, str)
+
         try:
-            value = self._store_dict[key]
+            value = self._store_dict[self.resolve_key(key)]
             start, length = _normalize_interval_index(value, byte_range)
             return prototype.buffer.from_buffer(value[start : start + length])
         except KeyError:
@@ -88,45 +89,46 @@ async def _get(key: str, byte_range: ByteRangeRequest) -> Buffer | None:
         return await concurrent_map(key_ranges, _get, limit=None)
 
     async def exists(self, key: str) -> bool:
-        return key in self._store_dict
+        return self.resolve_key(key) in self._store_dict
 
     async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None:
         self._check_writable()
         await self._ensure_open()
         assert isinstance(key, str)
         if not isinstance(value, Buffer):
             raise TypeError(f"Expected Buffer. Got {type(value)}.")
-
+        key_abs = self.resolve_key(key)
         if byte_range is not None:
-            buf = self._store_dict[key]
+            buf = self._store_dict[key_abs]
             buf[byte_range[0] : byte_range[1]] = value
-            self._store_dict[key] = buf
+            self._store_dict[key_abs] = buf
         else:
-            self._store_dict[key] = value
+            self._store_dict[key_abs] = value
 
     async def set_if_not_exists(self, key: str, default: Buffer) -> None:
         self._check_writable()
         await self._ensure_open()
-        self._store_dict.setdefault(key, default)
+        self._store_dict.setdefault(self.resolve_key(key), default)
 
     async def delete(self, key: str) -> None:
         self._check_writable()
         try:
-            del self._store_dict[key]
+            del self._store_dict[self.resolve_key(key)]
         except KeyError:
             pass  # Q(JH): why not raise?
 
     async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, bytes]]) -> None:
         raise NotImplementedError
 
     async def list(self) -> AsyncGenerator[str, None]:
-        for key in self._store_dict:
-            yield key
+        async for result in self.list_prefix(""):
+            yield result
 
     async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
+        prefix_abs = self.resolve_key(prefix)
         for key in self._store_dict:
-            if key.startswith(prefix):
-                yield key.removeprefix(prefix)
+            if key.startswith(prefix_abs):
+                yield key.removeprefix(prefix_abs).lstrip("/")
 
     async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
         """
@@ -141,8 +143,7 @@ async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
         -------
         AsyncGenerator[str, None]
         """
-        if prefix.endswith("/"):
-            prefix = prefix[:-1]
+        prefix = self.resolve_key(prefix)
 
         if prefix == "":
             keys_unique = {k.split("/")[0] for k in self._store_dict}

src/zarr/store/remote.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Self
+from typing import TYPE_CHECKING, Any
 
 import fsspec
 
@@ -10,6 +10,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Iterable
+    from typing import Self
 
     from fsspec.asyn import AsyncFileSystem