Skip to content

Commit

Permalink
progress on store interface (memory and local stores only)
Browse files Browse the repository at this point in the history
  • Loading branch information
jhamman committed Jan 16, 2024
1 parent f4822c7 commit ec2d43d
Show file tree
Hide file tree
Showing 13 changed files with 2,775 additions and 2,804 deletions.
4,964 changes: 2,485 additions & 2,479 deletions zarr/tests/test_storage.py

Large diffs are not rendered by default.

5 changes: 0 additions & 5 deletions zarr/v3/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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(
Expand Down
2 changes: 1 addition & 1 deletion zarr/v3/abc/codec.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
48 changes: 26 additions & 22 deletions zarr/v3/abc/store.py
Original file line number Diff line number Diff line change
@@ -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
-------
Expand All @@ -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
Expand All @@ -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.
Expand All @@ -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.
Expand All @@ -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
Expand All @@ -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:
Expand All @@ -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.
Expand Down Expand Up @@ -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
2 changes: 1 addition & 1 deletion zarr/v3/codecs/sharding.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@
CodecMetadata,
ShardingCodecIndexLocation,
)
from zarr.v3.stores import StorePath
from zarr.v3.store import StorePath

MAX_UINT_64 = 2**64 - 1

Expand Down
4 changes: 4 additions & 0 deletions zarr/v3/store/__init__.py
Original file line number Diff line number Diff line change
@@ -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
83 changes: 83 additions & 0 deletions zarr/v3/store/core.py
Original file line number Diff line number Diff line change
@@ -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
Loading

0 comments on commit ec2d43d

Please sign in to comment.