progress on store interface (memory and local stores only)

Author: jhamman

zarr/tests/test_storage.py

@@ -8,11 +8,7 @@
 from zarr.v3.group import Group  # noqa: F401
 from zarr.v3.metadata import RuntimeConfiguration, runtime_configuration  # noqa: F401
 from zarr.v3.store import (  # noqa: F401
-    LocalStore,
-    RemoteStore,
-    Store,
     StoreLike,
-    StorePath,
     make_store_path,
 )
 from zarr.v3.sync import sync as _sync
@@ -27,7 +23,6 @@ async def open_auto_async(
         return await Array.open(store_path, runtime_configuration=runtime_configuration_)
     except KeyError:
         return await Group.open(store_path, runtime_configuration=runtime_configuration_)
-    
 
 
 def open_auto(

zarr/v3/__init__.py

@@ -16,7 +16,7 @@
 import numpy as np
 
 from zarr.v3.common import BytesLike, SliceSelection
-from zarr.v3.stores import StorePath
+from zarr.v3.store import StorePath
 
 
 if TYPE_CHECKING:

zarr/v3/abc/codec.py

@@ -1,20 +1,19 @@
 from abc import abstractmethod, ABC
 
-from typing import List, Tuple
+from typing import List, Tuple, Optional
 
 
 class Store(ABC):
-    pass
-
-
-class ReadStore(Store):
     @abstractmethod
-    async def get(self, key: str) -> bytes:
+    async def get(
+        self, key: str, byte_range: Optional[Tuple[int, Optional[int]]] = None
+    ) -> Optional[bytes]:
         """Retrieve the value associated with a given key.
 
         Parameters
         ----------
         key : str
+        byte_range : tuple[int, Optional[int]], optional
 
         Returns
         -------
@@ -23,7 +22,9 @@ async def get(self, key: str) -> bytes:
         ...
 
     @abstractmethod
-    async def get_partial_values(self, key_ranges: List[Tuple[str, Tuple[int, int]]]) -> List[bytes]:
+    async def get_partial_values(
+        self, key_ranges: List[Tuple[str, Tuple[int, int]]]
+    ) -> List[bytes]:
         """Retrieve possibly partial values from given key_ranges.
 
         Parameters
@@ -38,6 +39,7 @@ async def get_partial_values(self, key_ranges: List[Tuple[str, Tuple[int, int]]]
         """
         ...
 
+    @abstractmethod
     async def exists(self, key: str) -> bool:
         """Check if a key exists in the store.
 
@@ -51,8 +53,12 @@ async def exists(self, key: str) -> bool:
         """
         ...
 
+    @property
+    @abstractmethod
+    def supports_writes(self) -> bool:
+        """Does the store support writes?"""
+        ...
 
-class WriteStore(ReadStore):
     @abstractmethod
     async def set(self, key: str, value: bytes) -> None:
         """Store a (key, value) pair.
@@ -64,7 +70,8 @@ async def set(self, key: str, value: bytes) -> None:
         """
         ...
 
-    async def delete(self, key: str) -> None
+    @abstractmethod
+    async def delete(self, key: str) -> None:
         """Remove a key from the store
 
         Parameters
@@ -73,10 +80,11 @@ async def delete(self, key: str) -> None
         """
         ...
 
-
-class PartialWriteStore(WriteStore):
-    # TODO, instead of using this, should we just check if the store is a PartialWriteStore?
-    supports_partial_writes = True 
+    @property
+    @abstractmethod
+    def supports_partial_writes(self) -> bool:
+        """Does the store support partial writes?"""
+        ...
 
     @abstractmethod
     async def set_partial_values(self, key_start_values: List[Tuple[str, int, bytes]]) -> None:
@@ -91,8 +99,12 @@ async def set_partial_values(self, key_start_values: List[Tuple[str, int, bytes]
         """
         ...
 
+    @property
+    @abstractmethod
+    def supports_listing(self) -> bool:
+        """Does the store support listing?"""
+        ...
 
-class ListMixin:
     @abstractmethod
     async def list(self) -> List[str]:
         """Retrieve all keys in the store.
@@ -132,11 +144,3 @@ async def list_dir(self, prefix: str) -> List[str]:
         list[str]
         """
         ...
-
-
-class ReadListStore(ReadStore, ListMixin):
-    pass
-
-
-class WriteListStore(WriteStore, ListMixin):
-    pass

zarr/v3/abc/store.py

@@ -43,7 +43,7 @@
     CodecMetadata,
     ShardingCodecIndexLocation,
 )
-from zarr.v3.stores import StorePath
+from zarr.v3.store import StorePath
 
 MAX_UINT_64 = 2**64 - 1
 

zarr/v3/codecs/sharding.py

@@ -0,0 +1,4 @@
+from zarr.v3.store.core import StorePath, StoreLike, make_store_path
+from zarr.v3.store.remote import RemoteStore
+from zarr.v3.store.local import LocalStore
+from zarr.v3.store.memory import MemoryStore

zarr/v3/store/__init__.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Optional, Tuple, Union
+
+from zarr.v3.common import BytesLike
+from zarr.v3.abc.store import Store
+
+
+def _dereference_path(root: str, path: str) -> str:
+    assert isinstance(root, str)
+    assert isinstance(path, str)
+    root = root.rstrip("/")
+    path = f"{root}/{path}" if root != "" else path
+    path = path.rstrip("/")
+    return path
+
+
+class StorePath:
+    store: Store
+    path: str
+
+    def __init__(self, store: Store, path: Optional[str] = None):
+        self.store = store
+        self.path = path or ""
+
+    @classmethod
+    def from_path(cls, pth: Path) -> StorePath:
+        return cls(Store.from_path(pth))
+
+    async def get(
+        self, byte_range: Optional[Tuple[int, Optional[int]]] = None
+    ) -> Optional[BytesLike]:
+        return await self.store.get(self.path, byte_range)
+
+    async def set(self, value: BytesLike, byte_range: Optional[Tuple[int, int]] = None) -> None:
+        if byte_range is not None:
+            raise NotImplementedError("Store.set does not have partial writes yet")
+        await self.store.set(self.path, value)
+
+    async def delete(self) -> None:
+        await self.store.delete(self.path)
+
+    async def exists(self) -> bool:
+        return await self.store.exists(self.path)
+
+    def __truediv__(self, other: str) -> StorePath:
+        return self.__class__(self.store, _dereference_path(self.path, other))
+
+    def __str__(self) -> str:
+        return _dereference_path(str(self.store), self.path)
+
+    def __repr__(self) -> str:
+        return f"StorePath({self.store.__class__.__name__}, {repr(str(self))})"
+
+    def __eq__(self, other: Any) -> bool:
+        try:
+            if self.store == other.store and self.path == other.path:
+                return True
+        except Exception:
+            pass
+        return False
+
+
+StoreLike = Union[Store, StorePath, Path, str]
+
+
+def make_store_path(store_like: StoreLike) -> StorePath:
+    if isinstance(store_like, StorePath):
+        return store_like
+    elif isinstance(store_like, Store):
+        return StorePath(store_like)
+    # elif isinstance(store_like, Path):
+    #     return StorePath(Store.from_path(store_like))
+    elif isinstance(store_like, str):
+        try:
+            from upath import UPath
+
+            return StorePath(Store.from_path(UPath(store_like)))
+        except ImportError as e:
+            raise e
+            # return StorePath(LocalStore(Path(store_like)))
+    raise TypeError