diff --git a/changes/3698.feature.md b/changes/3698.feature.md new file mode 100644 index 0000000000..d3a7ae0639 --- /dev/null +++ b/changes/3698.feature.md @@ -0,0 +1,8 @@ +`zarr.storage.ObjectStore` now accepts any object implementing the async +[obspec](https://developmentseed.org/obspec/) protocols, rather than only instances of +`obstore` store classes. The check is structural, so wrappers around an obstore store +(a cache, a request logger) can be used directly. Exceptions raised by the underlying +store are matched by obspec's well-known names, so a store that raises its own +`NotFoundError`, `NotSupportedError` or `AlreadyExistsError` gets the same handling as +obstore. The minimum supported `obstore` version is now 0.7.0 and the `remote` extra +also installs `obspec`. diff --git a/docs/user-guide/installation.md b/docs/user-guide/installation.md index a7487e83c8..9baca64d98 100644 --- a/docs/user-guide/installation.md +++ b/docs/user-guide/installation.md @@ -23,7 +23,7 @@ pip install zarr There are a number of optional dependency groups you can install for extra functionality. These can be installed using `pip install "zarr[]"`, e.g. `pip install "zarr[gpu]"` -- `remote`: support for reading/writing to remote data stores (fsspec, obstore) +- `remote`: support for reading/writing to remote data stores (fsspec, obstore, obspec) - `gpu`: support for GPUs (cupy) - `cli`: support for the `zarr` [command-line interface](cli.md) (typer) - `optional`: support for path-like access to local and remote stores (universal-pathlib) diff --git a/docs/user-guide/storage.md b/docs/user-guide/storage.md index 88ee2f1f32..b19952bdfc 100644 --- a/docs/user-guide/storage.md +++ b/docs/user-guide/storage.md @@ -193,9 +193,14 @@ print(array) ### Object Store -[`zarr.storage.ObjectStore`][] stores the contents of the Zarr hierarchy using any ObjectStore -[storage implementation](https://developmentseed.org/obstore/latest/api/store/), including AWS S3 ([`obstore.store.S3Store`][]), Google Cloud Storage ([`obstore.store.GCSStore`][]), and Azure Blob Storage ([`obstore.store.AzureStore`][]). This store is backed by [obstore](https://developmentseed.org/obstore/latest/), which -builds on the production quality Rust library [object_store](https://docs.rs/object_store/latest/object_store/). +[`zarr.storage.ObjectStore`][] stores the contents of the Zarr hierarchy using any object store +that implements the async [obspec](https://developmentseed.org/obspec/latest/) protocols. The +[obstore](https://developmentseed.org/obstore/latest/) stores are the usual choice: they cover AWS S3 +([`obstore.store.S3Store`][]), Google Cloud Storage ([`obstore.store.GCSStore`][]) and Azure Blob Storage +([`obstore.store.AzureStore`][]), and build on the production quality Rust library +[object_store](https://docs.rs/object_store/latest/object_store/). Because the requirement is structural +(the store only has to provide the obspec methods, not inherit from anything), middleware that wraps a +store, such as a cache or a request logger, works just as well. ```python exec="true" session="storage" source="above" result="ansi" from zarr.storage import ObjectStore diff --git a/pyproject.toml b/pyproject.toml index 4f5986edc7..13870961c0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -76,7 +76,8 @@ keywords = ["Python", "compressed", "ndimensional-arrays", "zarr"] # User-facing extras (shipped in package metadata) remote = [ "fsspec>=2023.10.0", - "obstore>=0.5.1", + "obstore>=0.7.0", + "obspec>=0.1.0", ] gpu = [ "cupy-cuda12x; sys_platform != 'darwin'", @@ -124,7 +125,8 @@ test = [ remote-tests = [ {include-group = "test"}, "fsspec>=2023.10.0", - "obstore>=0.5.1", + "obstore>=0.7.0", + "obspec>=0.1.0", "botocore", "s3fs>=2023.10.0", "moto[s3,server]==5.2.3", @@ -257,6 +259,7 @@ extra-dependencies = [ 'typing_extensions @ git+https://github.com/python/typing_extensions', 'donfig @ git+https://github.com/pytroll/donfig', 'obstore @ git+https://github.com/developmentseed/obstore@main#subdirectory=obstore', + 'obspec @ git+https://github.com/developmentseed/obspec@main', ] [tool.hatch.envs.upstream.env-vars] @@ -282,7 +285,8 @@ extra-dependencies = [ 'universal_pathlib==0.2.0', 'typing_extensions==4.14.*', 'donfig==0.8.*', - 'obstore==0.5.*', + 'obstore==0.7.*', + 'obspec==0.1.*', 'msgspec==0.19.*', ] diff --git a/src/zarr/__init__.py b/src/zarr/__init__.py index cdf3840c3b..2c0e32846d 100644 --- a/src/zarr/__init__.py +++ b/src/zarr/__init__.py @@ -79,6 +79,7 @@ def print_packages(packages: list[str]) -> None: "gcsfs", "universal-pathlib", "obstore", + "obspec", ] print(f"platform: {platform.platform()}") diff --git a/src/zarr/storage/_obstore.py b/src/zarr/storage/_obstore.py index b34a5f624d..4d5a97e812 100644 --- a/src/zarr/storage/_obstore.py +++ b/src/zarr/storage/_obstore.py @@ -1,12 +1,11 @@ from __future__ import annotations import asyncio -import contextlib import pickle from collections import defaultdict from itertools import chain from operator import itemgetter -from typing import TYPE_CHECKING, Self, TypedDict +from typing import TYPE_CHECKING, Literal, Self, TypedDict from zarr.abc.store import ( ByteRequest, @@ -21,30 +20,102 @@ if TYPE_CHECKING: from collections.abc import AsyncGenerator, Coroutine, Iterable, Sequence - from typing import Any - - from obstore import ListResult, ListStream, ObjectMeta, OffsetRange, SuffixRange - from obstore.store import ObjectStore as _UpstreamObjectStore + from collections.abc import Buffer as BufferLike + from typing import Any, Protocol + + from obspec import ( + DeleteAsync, + GetAsync, + GetRangeAsync, + GetRangesAsync, + HeadAsync, + ListAsync, + ListResult, + ListWithDelimiterAsync, + ObjectMeta, + OffsetRange, + PutAsync, + SuffixRange, + ) from zarr.core.buffer import Buffer, BufferPrototype + class ObspecInput( + DeleteAsync, + GetAsync, + GetRangeAsync, + GetRangesAsync, + HeadAsync, + ListAsync, + ListWithDelimiterAsync, + PutAsync, + Protocol, + ): + """The union of the async obspec protocols that ``ObjectStore`` relies on. + + Any object with these methods can back an ``ObjectStore``; there is no + requirement to inherit from anything. Keep ``_OBSPEC_METHODS`` in sync with + the protocols listed here. + """ + + __all__ = ["ObjectStore"] +# The methods of the ``ObspecInput`` protocol, checked structurally at runtime in +# ``ObjectStore.__init__``. obspec is only imported for type checking, so its +# protocol classes cannot be used for an ``isinstance`` check here. +_OBSPEC_METHODS: tuple[str, ...] = ( + "delete_async", + "get_async", + "get_range_async", + "get_ranges_async", + "head_async", + "list_async", + "list_with_delimiter_async", + "put_async", +) + _ALLOWED_EXCEPTIONS: tuple[type[Exception], ...] = ( FileNotFoundError, IsADirectoryError, NotADirectoryError, ) +_ObspecErrorName = Literal["AlreadyExistsError", "NotFoundError", "NotSupportedError"] + + +def _is_obspec_error(exc: Exception, name: _ObspecErrorName) -> bool: + """Check whether ``exc`` is the obspec exception called ``name``. + + obspec uses structural typing everywhere except for exceptions, which cannot be + matched structurally. Instead, implementations raise exceptions with well-known + class names and ``obspec.exceptions.map_exception`` resolves those names to the + obspec exception classes. The builtin ``FileNotFoundError`` maps to + ``NotFoundError``. -class ObjectStore[T_Store: "_UpstreamObjectStore"](Store): + obspec is imported lazily so that importing ``zarr.storage`` does not require it. """ - Store that uses obstore for fast read/write from AWS, GCP, Azure. + from obspec import exceptions + + return isinstance(exceptions.map_exception(exc), getattr(exceptions, name)) + + +class ObjectStore[T_Store: "ObspecInput"](Store): + """ + Store that reads and writes through any object store implementing the + [obspec](https://developmentseed.org/obspec/) async protocols, such as the + [obstore](https://developmentseed.org/obstore/) stores for AWS S3, Google + Cloud Storage and Azure Blob Storage, or a wrapper (a cache, a request logger) + around one of them. Parameters ---------- - store : obstore.store.ObjectStore - An obstore store instance that is set up with the proper credentials. + store : ObspecInput + Any object implementing the ``DeleteAsync``, ``GetAsync``, ``GetRangeAsync``, + ``GetRangesAsync``, ``HeadAsync``, ``ListAsync``, ``ListWithDelimiterAsync`` + and ``PutAsync`` obspec protocols, set up with the proper credentials. The + check is structural: the object must have those methods, but it does not + have to inherit from anything. read_only : bool Whether to open the store in read-only mode. @@ -55,7 +126,7 @@ class ObjectStore[T_Store: "_UpstreamObjectStore"](Store): """ store: T_Store - """The underlying obstore instance.""" + """The underlying obspec-compatible store instance.""" def __eq__(self, value: object) -> bool: if not isinstance(value, ObjectStore): @@ -67,8 +138,19 @@ def __eq__(self, value: object) -> bool: return self.store == value.store # type: ignore[no-any-return] def __init__(self, store: T_Store, *, read_only: bool = False) -> None: - if not store.__class__.__module__.startswith("obstore"): - raise TypeError(f"expected ObjectStore class, got {store!r}") + missing = [name for name in _OBSPEC_METHODS if not callable(getattr(store, name, None))] + if missing: + raise TypeError( + f"expected an object implementing the obspec async store protocols, got " + f"{store!r}, which is missing the following methods: {', '.join(missing)}" + ) + try: + import obspec # noqa: F401 + except ImportError as e: + raise ImportError( + "ObjectStore requires the obspec package. Install it with " + "'pip install obspec' or 'pip install zarr[remote]'." + ) from e super().__init__(read_only=read_only) self.store = store @@ -98,46 +180,31 @@ async def get( self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None ) -> Buffer | None: # docstring inherited - import obstore as obs - + if byte_range is not None and not isinstance( + byte_range, RangeByteRequest | OffsetByteRequest | SuffixByteRequest + ): + raise ValueError(f"Unexpected byte_range, got {byte_range}") try: if byte_range is None: - resp = await obs.get_async(self.store, key) - return prototype.buffer.from_bytes(await resp.bytes_async()) # type: ignore[arg-type] + resp = await self.store.get_async(key) + return prototype.buffer.from_bytes(await resp.buffer_async()) # type: ignore[arg-type] elif isinstance(byte_range, RangeByteRequest): - bytes = await obs.get_range_async( - self.store, key, start=byte_range.start, end=byte_range.end + bytes = await self.store.get_range_async( + key, start=byte_range.start, end=byte_range.end ) return prototype.buffer.from_bytes(bytes) # type: ignore[arg-type] elif isinstance(byte_range, OffsetByteRequest): - resp = await obs.get_async( - self.store, key, options={"range": {"offset": byte_range.offset}} + resp = await self.store.get_async( + key, options={"range": {"offset": byte_range.offset}} ) - return prototype.buffer.from_bytes(await resp.bytes_async()) # type: ignore[arg-type] - elif isinstance(byte_range, SuffixByteRequest): - # some object stores (Azure) don't support suffix requests. In this - # case, our workaround is to first get the length of the object and then - # manually request the byte range at the end. - try: - resp = await obs.get_async( - self.store, key, options={"range": {"suffix": byte_range.suffix}} - ) - return prototype.buffer.from_bytes(await resp.bytes_async()) # type: ignore[arg-type] - except obs.exceptions.NotSupportedError: - head_resp = await obs.head_async(self.store, key) - file_size = head_resp["size"] - suffix_len = byte_range.suffix - buffer = await obs.get_range_async( - self.store, - key, - start=file_size - suffix_len, - length=suffix_len, - ) - return prototype.buffer.from_bytes(buffer) # type: ignore[arg-type] + return prototype.buffer.from_bytes(await resp.buffer_async()) # type: ignore[arg-type] else: - raise ValueError(f"Unexpected byte_range, got {byte_range}") - except _ALLOWED_EXCEPTIONS: - return None + buffer = await _get_suffix(self.store, key, byte_range.suffix) + return prototype.buffer.from_bytes(buffer) # type: ignore[arg-type] + except Exception as e: + if isinstance(e, _ALLOWED_EXCEPTIONS) or _is_obspec_error(e, "NotFoundError"): + return None + raise async def get_partial_values( self, @@ -149,14 +216,13 @@ async def get_partial_values( async def exists(self, key: str) -> bool: # docstring inherited - import obstore as obs - try: - await obs.head_async(self.store, key) - except FileNotFoundError: - return False - else: - return True + await self.store.head_async(key) + except Exception as e: + if _is_obspec_error(e, "NotFoundError"): + return False + raise + return True @property def supports_writes(self) -> bool: @@ -165,21 +231,20 @@ def supports_writes(self) -> bool: async def set(self, key: str, value: Buffer) -> None: # docstring inherited - import obstore as obs - self._check_writable() buf = value.as_buffer_like() - await obs.put_async(self.store, key, buf) + await self.store.put_async(key, buf) async def set_if_not_exists(self, key: str, value: Buffer) -> None: # docstring inherited - import obstore as obs - self._check_writable() buf = value.as_buffer_like() - with contextlib.suppress(obs.exceptions.AlreadyExistsError): - await obs.put_async(self.store, key, buf, mode="create") + try: + await self.store.put_async(key, buf, mode="create") + except Exception as e: + if not _is_obspec_error(e, "AlreadyExistsError"): + raise @property def supports_deletes(self) -> bool: @@ -188,27 +253,25 @@ def supports_deletes(self) -> bool: async def delete(self, key: str) -> None: # docstring inherited - import obstore as obs - self._check_writable() - # Some obstore stores such as local filesystems, GCP and Azure raise an error + # Some stores such as local filesystems, GCP and Azure raise an error # when deleting a non-existent key, while others such as S3 and in-memory do - # not. We suppress the error to make the behavior consistent across all obstore + # not. We suppress the error to make the behavior consistent across all # stores. This is also in line with the behavior of the other Zarr store adapters. - with contextlib.suppress(FileNotFoundError): - await obs.delete_async(self.store, key) + try: + await self.store.delete_async(key) + except Exception as e: + if not _is_obspec_error(e, "NotFoundError"): + raise async def delete_dir(self, prefix: str) -> None: # docstring inherited - import obstore as obs - self._check_writable() if prefix != "" and not prefix.endswith("/"): prefix += "/" - metas = await obs.list(self.store, prefix).collect_async() - keys = [(m["path"],) for m in metas] + keys = [(obj["path"],) async for obj in self._list(prefix)] await concurrent_map(keys, self.delete, limit=config.get("async.concurrency")) @property @@ -217,10 +280,7 @@ def supports_listing(self) -> bool: return True async def _list(self, prefix: str | None = None) -> AsyncGenerator[ObjectMeta, None]: - import obstore as obs - - objects: ListStream[Sequence[ObjectMeta]] = obs.list(self.store, prefix=prefix) - async for batch in objects: + async for batch in self.store.list_async(prefix=prefix): for item in batch: yield item @@ -234,16 +294,19 @@ def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]: def list_dir(self, prefix: str) -> AsyncGenerator[str, None]: # docstring inherited - import obstore as obs - - coroutine = obs.list_with_delimiter_async(self.store, prefix=prefix) + coroutine = self.store.list_with_delimiter_async(prefix=prefix) return _transform_list_dir(coroutine, prefix) async def getsize(self, key: str) -> int: # docstring inherited - import obstore as obs - - resp = await obs.head_async(self.store, key) + try: + resp = await self.store.head_async(key) + except Exception as e: + # The Store contract is a FileNotFoundError; a store's own NotFoundError + # may not derive from it. + if not isinstance(e, FileNotFoundError) and _is_obspec_error(e, "NotFoundError"): + raise FileNotFoundError(key) from e + raise return resp["size"] async def getsize_prefix(self, prefix: str) -> int: @@ -252,6 +315,24 @@ async def getsize_prefix(self, prefix: str) -> int: return sum(sizes) +async def _get_suffix(store: ObspecInput, path: str, suffix: int) -> BufferLike: + """Fetch the last ``suffix`` bytes of ``path``. + + Some object stores (Azure) don't support suffix requests. In this case, our + workaround is to first get the length of the object and then manually request + the byte range at the end. + """ + try: + resp = await store.get_async(path, options={"range": {"suffix": suffix}}) + return await resp.buffer_async() + except Exception as e: + if not _is_obspec_error(e, "NotSupportedError"): + raise + head_resp = await store.head_async(path) + file_size = head_resp["size"] + return await store.get_range_async(path, start=file_size - suffix, length=suffix) + + async def _transform_list_dir( list_result_coroutine: Coroutine[Any, Any, ListResult[Sequence[ObjectMeta]]], prefix: str ) -> AsyncGenerator[str, None]: @@ -338,7 +419,7 @@ class _Response(TypedDict): async def _make_bounded_requests( - store: _UpstreamObjectStore, + store: ObspecInput, path: str, requests: list[_BoundedRequest], prototype: BufferPrototype, @@ -350,12 +431,10 @@ async def _make_bounded_requests( within a single file, and will e.g. merge concurrent requests. This only uses one single Python coroutine. """ - import obstore as obs - starts = [r["start"] for r in requests] ends = [r["end"] for r in requests] async with semaphore: - responses = await obs.get_ranges_async(store, path=path, starts=starts, ends=ends) + responses = await store.get_ranges_async(path=path, starts=starts, ends=ends) buffer_responses: list[_Response] = [] for request, response in zip(requests, responses, strict=True): @@ -370,7 +449,7 @@ async def _make_bounded_requests( async def _make_other_request( - store: _UpstreamObjectStore, + store: ObspecInput, request: _OtherRequest, prototype: BufferPrototype, semaphore: asyncio.Semaphore, @@ -380,14 +459,12 @@ async def _make_other_request( We return a `list[_Response]` for symmetry with `_make_bounded_requests` so that all futures can be gathered together. """ - import obstore as obs - async with semaphore: if request["range"] is None: - resp = await obs.get_async(store, request["path"]) + resp = await store.get_async(request["path"]) else: - resp = await obs.get_async(store, request["path"], options={"range": request["range"]}) - buffer = await resp.bytes_async() + resp = await store.get_async(request["path"], options={"range": request["range"]}) + buffer = await resp.buffer_async() return [ { @@ -398,7 +475,7 @@ async def _make_other_request( async def _make_suffix_request( - store: _UpstreamObjectStore, + store: ObspecInput, request: _SuffixRequest, prototype: BufferPrototype, semaphore: asyncio.Semaphore, @@ -406,28 +483,13 @@ async def _make_suffix_request( """Make suffix requests. This is separated out from `_make_other_request` because some object stores (Azure) - don't support suffix requests. In this case, our workaround is to first get the - length of the object and then manually request the byte range at the end. + don't support suffix requests; see `_get_suffix` for the workaround. We return a `list[_Response]` for symmetry with `_make_bounded_requests` so that all futures can be gathered together. """ - import obstore as obs - async with semaphore: - try: - resp = await obs.get_async(store, request["path"], options={"range": request["range"]}) - buffer = await resp.bytes_async() - except obs.exceptions.NotSupportedError: - head_resp = await obs.head_async(store, request["path"]) - file_size = head_resp["size"] - suffix_len = request["range"]["suffix"] - buffer = await obs.get_range_async( - store, - request["path"], - start=file_size - suffix_len, - length=suffix_len, - ) + buffer = await _get_suffix(store, request["path"], request["range"]["suffix"]) return [ { @@ -438,7 +500,7 @@ async def _make_suffix_request( async def _get_partial_values( - store: _UpstreamObjectStore, + store: ObspecInput, prototype: BufferPrototype, key_ranges: Iterable[tuple[str, ByteRequest | None]], ) -> list[Buffer | None]: diff --git a/tests/test_store/test_object.py b/tests/test_store/test_object.py index 1ea148b3c3..82cb467148 100644 --- a/tests/test_store/test_object.py +++ b/tests/test_store/test_object.py @@ -79,7 +79,7 @@ def test_store_equal(self, store: ObjectStore[LocalStore]) -> None: def test_store_init_raises(self) -> None: """Test __init__ raises appropriate error for improper store type""" - with pytest.raises(TypeError): + with pytest.raises(TypeError, match="missing the following methods: delete_async"): ObjectStore("path/to/store") # type: ignore[type-var] async def test_store_getsize(self, store: ObjectStore[LocalStore]) -> None: diff --git a/tests/test_store/test_object_obspec.py b/tests/test_store/test_object_obspec.py new file mode 100644 index 0000000000..ac63732686 --- /dev/null +++ b/tests/test_store/test_object_obspec.py @@ -0,0 +1,370 @@ +"""Tests for ``ObjectStore`` over a store that is not obstore. + +``ObjectStore`` accepts any object implementing the async obspec protocols. The store +defined here is a minimal pure-Python implementation of those protocols: it inherits +from nothing, and it raises its own exception classes, named after the obspec +exceptions but not derived from them or from the builtins, so these tests fail if +``ObjectStore`` ever falls back to nominal checks. +""" + +from __future__ import annotations + +import sys +from datetime import UTC, datetime +from typing import IO, TYPE_CHECKING, Any, TypedDict, cast + +import pytest + +pytest.importorskip("obspec") + +from zarr.abc.store import ByteRequest, OffsetByteRequest, RangeByteRequest, SuffixByteRequest +from zarr.core.buffer import Buffer, cpu, default_buffer_prototype +from zarr.storage import ObjectStore +from zarr.testing.store import StoreTests + +if TYPE_CHECKING: + from collections.abc import ( + AsyncIterable, + AsyncIterator, + Iterable, + Iterator, + Sequence, + ) + from collections.abc import ( + Buffer as BufferLike, + ) + from pathlib import Path + + from obspec import Attributes, GetOptions, ListResult, ObjectMeta, PutMode, PutResult + + +class NotFoundError(Exception): + """Deliberately not a ``FileNotFoundError``: only the name matches obspec.""" + + +class NotSupportedError(Exception): + pass + + +class AlreadyExistsError(Exception): + pass + + +class PermissionDeniedError(Exception): + """An obspec-named error that ``ObjectStore`` has no special handling for.""" + + +class _GetResult: + def __init__(self, path: str, data: bytes, byte_range: tuple[int, int]) -> None: + self._path = path + self._data = data + self._range = byte_range + + @property + def attributes(self) -> Attributes: + return {} + + @property + def meta(self) -> ObjectMeta: + return _meta(self._path, self._data) + + @property + def range(self) -> tuple[int, int]: + return self._range + + async def buffer_async(self) -> BufferLike: + return self._data + + async def __aiter__(self) -> AsyncIterator[BufferLike]: + yield self._data + + +def _meta(path: str, data: bytes) -> ObjectMeta: + return { + "path": path, + "last_modified": datetime.now(UTC), + "size": len(data), + "e_tag": None, + "version": None, + } + + +class DictObspecStore: + """An in-memory implementation of the async obspec protocols ``ObjectStore`` uses. + + Suffix range requests raise ``NotSupportedError`` (as Azure does) unless + ``supports_suffix`` is set, so that ``ObjectStore`` has to take its fallback path. + """ + + def __init__(self, *, supports_suffix: bool = False) -> None: + self.data: dict[str, bytes] = {} + self.supports_suffix = supports_suffix + + def __eq__(self, other: object) -> bool: + return ( + isinstance(other, DictObspecStore) + and other.data == self.data + and other.supports_suffix == self.supports_suffix + ) + + def _read(self, path: str) -> bytes: + try: + return self.data[path] + except KeyError: + raise NotFoundError(path) from None + + async def head_async(self, path: str) -> ObjectMeta: + return _meta(path, self._read(path)) + + async def get_async(self, path: str, *, options: GetOptions | None = None) -> _GetResult: + data = self._read(path) + byte_range = (options or {}).get("range") + if byte_range is None: + start, end = 0, len(data) + elif isinstance(byte_range, dict): + range_dict = cast("dict[str, int]", byte_range) + if "suffix" in range_dict: + if not self.supports_suffix: + raise NotSupportedError("suffix range requests are not supported") + start, end = max(len(data) - range_dict["suffix"], 0), len(data) + else: + start, end = range_dict["offset"], len(data) + else: + start, end = byte_range[0], byte_range[1] + return _GetResult(path, data[start:end], (start, end)) + + async def get_range_async( + self, path: str, *, start: int, end: int | None = None, length: int | None = None + ) -> BufferLike: + if end is None: + assert length is not None + end = start + length + return self._read(path)[start:end] + + async def get_ranges_async( + self, + path: str, + *, + starts: Sequence[int], + ends: Sequence[int] | None = None, + lengths: Sequence[int] | None = None, + ) -> Sequence[BufferLike]: + if ends is None: + assert lengths is not None + ends = [start + length for start, length in zip(starts, lengths, strict=True)] + data = self._read(path) + return [data[start:end] for start, end in zip(starts, ends, strict=True)] + + async def put_async( + self, + path: str, + file: IO[bytes] + | Path + | bytes + | BufferLike + | AsyncIterator[BufferLike] + | AsyncIterable[BufferLike] + | Iterator[BufferLike] + | Iterable[BufferLike], + *, + attributes: Attributes | None = None, + tags: dict[str, str] | None = None, + mode: PutMode | None = None, + use_multipart: bool | None = None, + chunk_size: int = 5 * 1024 * 1024, + max_concurrency: int = 12, + ) -> PutResult: + if mode == "create" and path in self.data: + raise AlreadyExistsError(path) + self.data[path] = bytes(memoryview(file)) # type: ignore[arg-type] + return {"e_tag": None, "version": None} + + async def delete_async(self, paths: str | Sequence[str]) -> None: + for path in [paths] if isinstance(paths, str) else paths: + if path not in self.data: + raise NotFoundError(path) + del self.data[path] + + async def list_async( + self, prefix: str | None = None, *, offset: str | None = None + ) -> AsyncIterator[Sequence[ObjectMeta]]: + # obspec evaluates prefixes per path segment: "c" matches "c/0" but not "cc/0". + base = (prefix or "").strip("/") + yield [ + _meta(k, v) + for k, v in sorted(self.data.items()) + if not base or k == base or k.startswith(base + "/") + ] + + async def list_with_delimiter_async( + self, prefix: str | None = None + ) -> ListResult[Sequence[ObjectMeta]]: + base = (prefix or "").strip("/") + base = base + "/" if base else "" + common_prefixes: set[str] = set() + objects: list[ObjectMeta] = [] + for key, value in sorted(self.data.items()): + if not key.startswith(base): + continue + child, sep, _ = key[len(base) :].partition("/") + if sep: + common_prefixes.add(base + child) + else: + objects.append(_meta(key, value)) + return {"common_prefixes": sorted(common_prefixes), "objects": objects} + + +class StoreKwargs(TypedDict): + store: DictObspecStore + read_only: bool + + +class TestObspecObjectStore(StoreTests[ObjectStore[DictObspecStore], cpu.Buffer]): + # store_cls is needed to do an isinstance check, so can't be a subscripted generic + store_cls = ObjectStore # type: ignore[assignment] + buffer_cls = cpu.Buffer + + @pytest.fixture + def store_kwargs(self) -> StoreKwargs: + return {"store": DictObspecStore(), "read_only": False} + + @pytest.fixture + def store(self, store_kwargs: StoreKwargs) -> ObjectStore[DictObspecStore]: + return self.store_cls(**store_kwargs) + + async def get(self, store: ObjectStore[DictObspecStore], key: str) -> Buffer: + return self.buffer_cls.from_bytes(store.store.data[key]) + + async def set(self, store: ObjectStore[DictObspecStore], key: str, value: Buffer) -> None: + store.store.data[key] = value.to_bytes() + + def test_store_repr(self, store: ObjectStore[DictObspecStore]) -> None: + assert repr(store).startswith("ObjectStore(object_store://") + + def test_store_supports_writes(self, store: ObjectStore[DictObspecStore]) -> None: + assert store.supports_writes + + def test_store_supports_partial_writes(self, store: ObjectStore[DictObspecStore]) -> None: + assert not store.supports_partial_writes + + def test_store_supports_listing(self, store: ObjectStore[DictObspecStore]) -> None: + assert store.supports_listing + + +@pytest.mark.parametrize("supports_suffix", [True, False]) +async def test_suffix_requests(supports_suffix: bool) -> None: + """Suffix reads give the same result whether or not the store supports them natively.""" + store = ObjectStore(DictObspecStore(supports_suffix=supports_suffix)) + await store.set("key", cpu.Buffer.from_bytes(b"0123456789")) + prototype = default_buffer_prototype() + + single = await store.get("key", prototype, SuffixByteRequest(3)) + assert single is not None + assert single.to_bytes() == b"789" + + ranges: list[tuple[str, ByteRequest | None]] = [ + ("key", SuffixByteRequest(2)), + ("key", RangeByteRequest(1, 3)), + ("key", OffsetByteRequest(8)), + ("key", None), + ] + observed = await store.get_partial_values(prototype, ranges) + assert [buf.to_bytes() for buf in observed if buf is not None] == [ + b"89", + b"12", + b"89", + b"0123456789", + ] + + +async def test_set_if_not_exists_keeps_existing_value() -> None: + store = ObjectStore(DictObspecStore()) + await store.set("key", cpu.Buffer.from_bytes(b"first")) + await store.set_if_not_exists("key", cpu.Buffer.from_bytes(b"second")) + assert store.store.data["key"] == b"first" + + +async def test_wrapping_an_obstore_store() -> None: + """A wrapper that delegates to an obstore store is accepted, not just obstore itself.""" + pytest.importorskip("obstore") + from obstore.store import MemoryStore + + class LoggingStore: + def __init__(self, inner: MemoryStore) -> None: + self.inner = inner + self.calls: list[str] = [] + + def __getattr__(self, name: str) -> Any: + attr = getattr(self.inner, name) + if callable(attr): + self.calls.append(name) + return attr + + wrapper = LoggingStore(MemoryStore()) + store = ObjectStore(wrapper) + await store.set("a", cpu.Buffer.from_bytes(b"x")) + result = await store.get("a", default_buffer_prototype()) + assert result is not None + assert result.to_bytes() == b"x" + assert not await store.exists("b") + assert "put_async" in wrapper.calls + assert "head_async" in wrapper.calls + + +def test_init_rejects_object_without_obspec_methods() -> None: + class PartialStore: + async def get_async(self, path: str) -> None: ... + + with pytest.raises(TypeError, match=r"missing the following methods: delete_async, "): + ObjectStore(PartialStore()) # type: ignore[type-var] + + +def test_init_requires_obspec(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem(sys.modules, "obspec", None) + with pytest.raises(ImportError, match="ObjectStore requires the obspec package"): + ObjectStore(DictObspecStore()) + + +async def test_unhandled_store_errors_propagate_unchanged() -> None: + """Errors that ObjectStore does not handle are re-raised as-is, not as obspec copies.""" + store = ObjectStore(DictObspecStore()) + error = PermissionDeniedError("nope") + + async def deny(*args: Any, **kwargs: Any) -> Any: + raise error + + store.store.head_async = deny # type: ignore[method-assign] + store.store.get_async = deny # type: ignore[method-assign] + store.store.delete_async = deny # type: ignore[method-assign] + store.store.put_async = deny # type: ignore[method-assign] + + with pytest.raises(PermissionDeniedError) as info: + await store.exists("key") + assert info.value is error + with pytest.raises(PermissionDeniedError) as info: + await store.get("key", default_buffer_prototype()) + assert info.value is error + with pytest.raises(PermissionDeniedError) as info: + await store.delete("key") + assert info.value is error + with pytest.raises(PermissionDeniedError) as info: + await store.set_if_not_exists("key", cpu.Buffer.from_bytes(b"")) + assert info.value is error + + +@pytest.mark.parametrize( + "byte_range", [None, RangeByteRequest(0, 2), OffsetByteRequest(1), SuffixByteRequest(1)] +) +async def test_get_missing_key_returns_none(byte_range: ByteRequest | None) -> None: + """A store's own NotFoundError is treated like the builtin FileNotFoundError.""" + store = ObjectStore(DictObspecStore()) + assert await store.get("missing", default_buffer_prototype(), byte_range) is None + + +async def test_get_partial_values_rejects_unknown_range() -> None: + store = ObjectStore(DictObspecStore()) + with pytest.raises(ValueError, match="Unsupported range input"): + await store.get_partial_values( + default_buffer_prototype(), + [("key", (0, 1))], # type: ignore[list-item] + ) diff --git a/uv.lock b/uv.lock index f622280eb7..cef8659a48 100644 --- a/uv.lock +++ b/uv.lock @@ -2056,6 +2056,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/62/5e/3a6a3e90f35cea3853c45e5d5fb9b7192ce4384616f932cf7591298ab6e1/numpydoc-1.10.0-py3-none-any.whl", hash = "sha256:3149da9874af890bcc2a82ef7aae5484e5aa81cb2778f08e3c307ba6d963721b", size = 69255, upload-time = "2025-12-02T16:39:11.561Z" }, ] +[[package]] +name = "obspec" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e2/94/7a9ad6927cac6ec7680e11772fb692145a05a93bafd80b84f6f0ef12f4e7/obspec-0.1.0.tar.gz", hash = "sha256:b189781a53f82ef8d6abf0c9e77fd4c46ac9f244d5a91eb35ee61c2e2b204a4a", size = 117254, upload-time = "2025-06-25T05:24:00.002Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bb/69/96feeac84ce0b871567225c78515f3b557c023e72ed9b4f1833f3662bd6b/obspec-0.1.0-py3-none-any.whl", hash = "sha256:307f0fa2c2998b324ecf0eed6a2a89049a4c40c9b1fa2b5e1af28f0ee72136b3", size = 15231, upload-time = "2025-06-25T05:23:58.735Z" }, +] + [[package]] name = "obstore" version = "0.11.1" @@ -3581,6 +3590,7 @@ optional = [ ] remote = [ { name = "fsspec" }, + { name = "obspec" }, { name = "obstore" }, ] @@ -3603,6 +3613,7 @@ dev = [ { name = "mypy" }, { name = "numcodecs", extra = ["msgpack"] }, { name = "numpydoc" }, + { name = "obspec" }, { name = "obstore" }, { name = "pytest" }, { name = "pytest-accept" }, @@ -3646,6 +3657,7 @@ remote-tests = [ { name = "hypothesis" }, { name = "moto", extra = ["s3", "server"] }, { name = "numpydoc" }, + { name = "obspec" }, { name = "obstore" }, { name = "pytest" }, { name = "pytest-accept" }, @@ -3686,7 +3698,8 @@ requires-dist = [ { name = "msgspec", specifier = ">=0.19" }, { name = "numcodecs", specifier = ">=0.14" }, { name = "numpy", specifier = ">=2" }, - { name = "obstore", marker = "extra == 'remote'", specifier = ">=0.5.1" }, + { name = "obspec", marker = "extra == 'remote'", specifier = ">=0.1.0" }, + { name = "obstore", marker = "extra == 'remote'", specifier = ">=0.7.0" }, { name = "packaging", specifier = ">=22.0" }, { name = "typer", marker = "extra == 'cli'" }, { name = "typing-extensions", specifier = ">=4.14" }, @@ -3713,7 +3726,8 @@ dev = [ { name = "mypy", specifier = "==2.3.1" }, { name = "numcodecs", extras = ["msgpack"] }, { name = "numpydoc", specifier = "==1.10.0" }, - { name = "obstore", specifier = ">=0.5.1" }, + { name = "obspec", specifier = ">=0.1.0" }, + { name = "obstore", specifier = ">=0.7.0" }, { name = "pytest", specifier = "==9.1.1" }, { name = "pytest-accept", specifier = "==0.3.0" }, { name = "pytest-asyncio", specifier = "==1.4.0" }, @@ -3754,7 +3768,8 @@ remote-tests = [ { name = "hypothesis", specifier = "==6.165.10" }, { name = "moto", extras = ["s3", "server"], specifier = "==5.2.3" }, { name = "numpydoc", specifier = "==1.10.0" }, - { name = "obstore", specifier = ">=0.5.1" }, + { name = "obspec", specifier = ">=0.1.0" }, + { name = "obstore", specifier = ">=0.7.0" }, { name = "pytest", specifier = "==9.1.1" }, { name = "pytest-accept", specifier = "==0.3.0" }, { name = "pytest-asyncio", specifier = "==1.4.0" },