diff --git a/.gitignore b/.gitignore index 8556e70..35540fc 100644 --- a/.gitignore +++ b/.gitignore @@ -49,12 +49,6 @@ coverage.xml # Django stuff: *.log -# Sphinx documentation -docs/_build/ -docs/api -docs/data -data -data.zip # PyBuilder target/ @@ -84,4 +78,7 @@ junit.xml tests/.hypothesis .hypothesis/ -site/* \ No newline at end of file +site/* + +# uv lock files +uv.lock diff --git a/docs/api.md b/docs/api.md deleted file mode 100644 index 0f60b9f..0000000 --- a/docs/api.md +++ /dev/null @@ -1,12 +0,0 @@ -# API - -::: obspec_utils - -## Typing - -::: obspec_utils.typing.Url -::: obspec_utils.typing.Path - -## Registry Utilities - -::: obspec_utils.registry.UrlKey diff --git a/docs/api/aiohttp.md b/docs/api/aiohttp.md new file mode 100644 index 0000000..9631889 --- /dev/null +++ b/docs/api/aiohttp.md @@ -0,0 +1 @@ +::: obspec_utils.aiohttp.AiohttpStore diff --git a/docs/api/obspec.md b/docs/api/obspec.md new file mode 100644 index 0000000..5ae5359 --- /dev/null +++ b/docs/api/obspec.md @@ -0,0 +1,2 @@ +::: obspec_utils.obspec.StoreReader +::: obspec_utils.obspec.StoreMemCacheReader diff --git a/docs/api/obstore.md b/docs/api/obstore.md new file mode 100644 index 0000000..d31868e --- /dev/null +++ b/docs/api/obstore.md @@ -0,0 +1,2 @@ +::: obspec_utils.obstore.ObstoreReader +::: obspec_utils.obstore.ObstoreMemCacheReader diff --git a/docs/api/registry.md b/docs/api/registry.md new file mode 100644 index 0000000..4812ac0 --- /dev/null +++ b/docs/api/registry.md @@ -0,0 +1,2 @@ +::: obspec_utils.registry.ObjectStoreRegistry +::: obspec_utils.registry.UrlKey diff --git a/docs/api/typing.md b/docs/api/typing.md new file mode 100644 index 0000000..f6ab911 --- /dev/null +++ b/docs/api/typing.md @@ -0,0 +1,4 @@ +::: obspec_utils.obspec.ReadableStore + +::: obspec_utils.typing.Url +::: obspec_utils.typing.Path diff --git a/docs/index.md b/docs/index.md index 09a5d83..3b7e795 100644 --- a/docs/index.md +++ b/docs/index.md @@ -8,7 +8,23 @@ Utilities for interacting with object storage, based on [obspec](https://github. 1. **ObjectStoreRegistry**: A registry for managing multiple object stores, allowing you to resolve URLs to the appropriate store and path. This is particularly useful when working with datasets that span multiple storage backends or buckets. -2. **File Handlers**: Wrappers around obstore's file reading capabilities that provide a familiar file-like interface, making it easy to integrate with libraries that expect standard Python file objects. +2. **ReadableStore Protocol**: A minimal protocol defining the read-only interface required for object storage access. This allows alternative backends (like aiohttp) to be used instead of obstore. + +3. **File Handlers**: Wrappers around obstore's file reading capabilities that provide a familiar file-like interface. + +## Design Philosophy + +The library is designed around **protocols rather than concrete classes**. The `ObjectStoreRegistry` accepts any object that implements the `ReadableStore` protocol, which means: + +- **obstore classes** (S3Store, HTTPStore, GCSStore, etc.) work out of the box +- **Custom implementations** (like the included `AiohttpStore`) can be used as alternatives +- **The Zarr/VirtualiZarr layer doesn't care** which backend you use - it just needs something satisfying the protocol + +This is particularly useful when: + +- obstore's HTTPStore (designed for WebDAV/S3-like semantics) isn't ideal for your use case +- You need generic HTTPS access to THREDDS, NASA data servers, or other HTTP endpoints +- You want to use a different HTTP library like aiohttp ## Getting started @@ -26,7 +42,7 @@ The `ObjectStoreRegistry` allows you to register object stores and resolve URLs ```python from obstore.store import S3Store -from obspec_utils import ObjectStoreRegistry +from obspec_utils.registry import ObjectStoreRegistry # Create and register stores s3store = S3Store(bucket="my-bucket", prefix="my-data/") @@ -37,23 +53,75 @@ store, path = registry.resolve("s3://my-bucket/my-data/file.nc") # path == "file.nc" ``` +### Using Alternative HTTP Backends + +For generic HTTPS access where obstore's HTTPStore may not be ideal, you can use the `AiohttpStore`: + +```python +from obspec_utils.registry import ObjectStoreRegistry +from obspec_utils.aiohttp import AiohttpStore + +# Create an aiohttp-based store for a THREDDS server +store = AiohttpStore( + "https://thredds.example.com/data", + headers={"Authorization": "Bearer "}, # Optional auth + timeout=60.0, +) + +registry = ObjectStoreRegistry({"https://thredds.example.com/data": store}) + +# Use it just like any other store +store, path = registry.resolve("https://thredds.example.com/data/file.nc") +data = await store.get_range_async(path, start=0, end=1000) +``` + ### File Handlers -The file handlers provide file-like interfaces for reading from object stores: +The file handlers provide file-like interfaces (read, seek, tell) for reading from object stores. + +#### Protocol-based readers (recommended) + +These work with **any** ReadableStore implementation: ```python -from obstore.store import S3Store -from obspec_utils import ObstoreReader, ObstoreMemCacheReader +from obspec_utils.obspec import StoreReader, StoreMemCacheReader +# Works with obstore +from obstore.store import S3Store store = S3Store(bucket="my-bucket") +reader = StoreReader(store, "path/to/file.bin", buffer_size=1024*1024) + +# Also works with AiohttpStore or any ReadableStore +from obspec_utils.aiohttp import AiohttpStore +store = AiohttpStore("https://example.com/data") +reader = StoreReader(store, "file.bin") # Standard reader with buffered reads -reader = ObstoreReader(store, "path/to/file.bin", buffer_size=1024*1024) data = reader.read(100) # Read 100 bytes +reader.seek(0) # Seek back to start # Memory-cached reader for repeated access +cached_reader = StoreMemCacheReader(store, "file.bin") +data = cached_reader.readall() +``` + +#### Obstore-specific readers + +For maximum performance with obstore, use the obstore-specific readers which leverage obstore's native `ReadableFile`: + +```python +from obstore.store import S3Store +from obspec_utils.obstore import ObstoreReader, ObstoreMemCacheReader + +store = S3Store(bucket="my-bucket") + +# Uses obstore's optimized buffered reader +reader = ObstoreReader(store, "path/to/file.bin", buffer_size=1024*1024) +data = reader.read(100) + +# Uses obstore's MemoryStore for caching cached_reader = ObstoreMemCacheReader(store, "path/to/file.bin") -data = cached_reader.readall() # Read entire file from memory cache +data = cached_reader.readall() ``` ## Contributing diff --git a/mkdocs.yml b/mkdocs.yml index 204a1fb..c7c60d1 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -14,7 +14,12 @@ extra: nav: - "index.md" - - "api.md" + - "API": + - Typing: "api/typing.md" + - Aiohttp Adapters: "api/aiohttp.md" + - Obspec File Readers: "api/obspec.md" + - Obstore File Readers: "api/obstore.md" + - Store Registries: "api/registry.md" watch: - src/obspec_utils diff --git a/src/obspec_utils/__init__.py b/src/obspec_utils/__init__.py index f0de586..26d23ba 100644 --- a/src/obspec_utils/__init__.py +++ b/src/obspec_utils/__init__.py @@ -1,10 +1,3 @@ from ._version import __version__ -from .file_handlers import ObstoreMemCacheReader, ObstoreReader -from .registry import ObjectStoreRegistry -__all__ = [ - "__version__", - "ObstoreMemCacheReader", - "ObstoreReader", - "ObjectStoreRegistry", -] +__all__ = ["__version__"] diff --git a/src/obspec_utils/aiohttp.py b/src/obspec_utils/aiohttp.py new file mode 100644 index 0000000..86e628a --- /dev/null +++ b/src/obspec_utils/aiohttp.py @@ -0,0 +1,558 @@ +""" +Aiohttp-based implementation of the ReadableStore protocol. + +This module provides an alternative HTTP backend using aiohttp instead of obstore's HTTPStore. +It's useful for generic HTTPS access (e.g., THREDDS, NASA data from outside AWS region) +where obstore's HTTPStore (designed for WebDAV/S3-like semantics) may not be ideal. + +Example +------- + +```python +from obspec_utils.registry import ObjectStoreRegistry +from obspec_utils.aiohttp import AiohttpStore + +# Use the store as an async context manager for efficient session reuse +async with AiohttpStore("https://example.com/data") as store: + # Register it with the registry + registry = ObjectStoreRegistry({"https://example.com/data": store}) + + # Now VirtualiZarr can use this store for HTTP access + resolved_store, path = registry.resolve("https://example.com/data/file.nc") + data = await resolved_store.get_range_async(path, start=0, end=1000) +``` +""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator, Iterator, Sequence +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import TYPE_CHECKING + +from obspec import GetResult, GetResultAsync + +from obspec_utils.obspec import ReadableStore + +if TYPE_CHECKING: + from obspec import Attributes, GetOptions, ObjectMeta + +try: + import aiohttp +except ImportError as e: + raise ImportError( + "aiohttp is required for AiohttpStore. Install it with: pip install aiohttp" + ) from e + + +@dataclass +class AiohttpGetResult(GetResult): + """ + Result from a get request using aiohttp. + + Implements the obspec GetResult protocol for synchronous iteration. + """ + + _data: bytes + _meta: ObjectMeta + _attributes: Attributes = field(default_factory=dict) + _range: tuple[int, int] = (0, 0) + + def __post_init__(self): + if self._range == (0, 0): + self._range = (0, len(self._data)) + + @property + def attributes(self) -> Attributes: + """Additional object attributes.""" + return self._attributes + + def buffer(self) -> bytes: + """Return the data as a buffer.""" + return self._data + + @property + def meta(self) -> ObjectMeta: + """The ObjectMeta for this object.""" + return self._meta + + @property + def range(self) -> tuple[int, int]: + """The range of bytes returned by this request.""" + return self._range + + def __iter__(self) -> Iterator[bytes]: + """Iterate over chunks of the data.""" + yield self._data + + +@dataclass +class AiohttpGetResultAsync(GetResultAsync): + """ + Result from an async get request using aiohttp. + + Implements the obspec GetResultAsync protocol for asynchronous iteration. + """ + + _data: bytes + _meta: ObjectMeta + _attributes: Attributes = field(default_factory=dict) + _range: tuple[int, int] = (0, 0) + + def __post_init__(self): + if self._range == (0, 0): + self._range = (0, len(self._data)) + + @property + def attributes(self) -> Attributes: + """Additional object attributes.""" + return self._attributes + + async def buffer_async(self) -> bytes: + """Return the data as a buffer.""" + return self._data + + @property + def meta(self) -> ObjectMeta: + """The ObjectMeta for this object.""" + return self._meta + + @property + def range(self) -> tuple[int, int]: + """The range of bytes returned by this request.""" + return self._range + + async def __aiter__(self) -> AsyncIterator[bytes]: + """Async iterate over chunks of the data.""" + yield self._data + + +class AiohttpStore(ReadableStore): + """ + An aiohttp-based implementation of the ReadableStore protocol. + + This provides a lightweight alternative to obstore's HTTPStore for generic + HTTP/HTTPS access. It's particularly useful for: + + - THREDDS data servers + - NASA data access from outside AWS regions + - Any generic HTTP endpoint that doesn't need S3-like semantics + + The store should be used as an async context manager to efficiently reuse + a single HTTP session across multiple requests. + + Parameters + ---------- + base_url + The base URL for this store. All paths are resolved relative to this URL. + headers + Optional HTTP headers to include in all requests (e.g., authentication). + timeout + Request timeout in seconds. Default is 30. + + Examples + -------- + + Recommended usage with async context manager: + + ```python + async with AiohttpStore("https://example.com/data") as store: + # All requests share the same session + result = await store.get_async("file.nc") + data = await result.buffer_async() + + # Byte range requests + chunk = await store.get_range_async("file.nc", start=0, end=1000) + ``` + + Synchronous usage (creates a session per request): + + ```python + store = AiohttpStore("https://example.com/data") + result = store.get("file.nc") + data = result.buffer() + ``` + + With authentication: + + ```python + async with AiohttpStore( + "https://api.example.com/data", + headers={"Authorization": "Bearer "} + ) as store: + result = await store.get_async("protected/file.nc") + ``` + """ + + def __init__( + self, + base_url: str, + *, + headers: dict[str, str] | None = None, + timeout: float = 30.0, + ) -> None: + self.base_url = base_url.rstrip("/") + self.headers = headers or {} + self.timeout = aiohttp.ClientTimeout(total=timeout) + self._session: aiohttp.ClientSession | None = None + + async def __aenter__(self) -> "AiohttpStore": + """Enter the async context manager, creating a reusable session.""" + self._session = aiohttp.ClientSession( + timeout=self.timeout, + headers=self.headers, + ) + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: + """Exit the async context manager, closing the session.""" + if self._session is not None: + await self._session.close() + self._session = None + + def _build_url(self, path: str) -> str: + """Build the full URL from base URL and path.""" + path = path.removeprefix("/") + return f"{self.base_url}/{path}" if path else self.base_url + + def _parse_meta_from_headers( + self, path: str, headers: dict, content_length: int | None = None + ) -> ObjectMeta: + """Extract ObjectMeta from HTTP response headers.""" + # Parse last-modified header + last_modified_str = headers.get("Last-Modified") + if last_modified_str: + # Parse HTTP date format + try: + from email.utils import parsedate_to_datetime + + last_modified = parsedate_to_datetime(last_modified_str) + except (ValueError, TypeError): + last_modified = datetime.now(timezone.utc) + else: + last_modified = datetime.now(timezone.utc) + + # Get size from Content-Length or Content-Range + size = content_length or 0 + content_range = headers.get("Content-Range") + if content_range and "/" in content_range: + # Format: bytes 0-999/1234 + total_str = content_range.split("/")[-1] + if total_str != "*": + size = int(total_str) + elif "Content-Length" in headers: + size = int(headers["Content-Length"]) + + return { + "path": path, + "last_modified": last_modified, + "size": size, + "e_tag": headers.get("ETag"), + "version": None, + } + + def _parse_attributes_from_headers(self, headers: dict) -> Attributes: + """Extract Attributes from HTTP response headers.""" + attrs: Attributes = {} + header_map = { + "Content-Disposition": "Content-Disposition", + "Content-Encoding": "Content-Encoding", + "Content-Language": "Content-Language", + "Content-Type": "Content-Type", + "Cache-Control": "Cache-Control", + } + for attr_name, header_name in header_map.items(): + if header_name in headers: + attrs[attr_name] = headers[header_name] + return attrs + + # --- Async methods (primary implementation) --- + + async def _do_get_async( + self, + session: aiohttp.ClientSession, + path: str, + *, + options: GetOptions | None = None, + ) -> AiohttpGetResultAsync: + """Internal method that performs the actual GET request.""" + url = self._build_url(path) + request_headers = {} if self._session else dict(self.headers) + + # Handle range option if specified + byte_range = (0, 0) + if options and "range" in options: + range_opt = options["range"] + if isinstance(range_opt, tuple): + start, end = range_opt[0], range_opt[1] + request_headers["Range"] = f"bytes={start}-{end - 1}" + byte_range = (start, end) + elif isinstance(range_opt, dict): + if "offset" in range_opt: + request_headers["Range"] = f"bytes={range_opt['offset']}-" + elif "suffix" in range_opt: + request_headers["Range"] = f"bytes=-{range_opt['suffix']}" + + async with session.get(url, headers=request_headers) as response: + response.raise_for_status() + data = await response.read() + meta = self._parse_meta_from_headers( + path, dict(response.headers), len(data) + ) + attrs = self._parse_attributes_from_headers(dict(response.headers)) + + if byte_range == (0, 0): + byte_range = (0, len(data)) + + return AiohttpGetResultAsync( + _data=data, + _meta=meta, + _attributes=attrs, + _range=byte_range, + ) + + async def get_async( + self, + path: str, + *, + options: GetOptions | None = None, + ) -> AiohttpGetResultAsync: + """ + Download a file asynchronously. + + Parameters + ---------- + path + Path to the file relative to base_url. + options + Optional get options (range, conditionals, etc.). + + Returns + ------- + AiohttpGetResultAsync + Result object with buffer_async() method and metadata. + """ + if self._session is not None: + return await self._do_get_async(self._session, path, options=options) + + # Fallback: create a temporary session for this request + async with aiohttp.ClientSession( + timeout=self.timeout, headers=self.headers + ) as session: + return await self._do_get_async(session, path, options=options) + + async def _do_get_range_async( + self, + session: aiohttp.ClientSession, + path: str, + *, + start: int, + end: int, + ) -> bytes: + """Internal method that performs the actual range GET request.""" + url = self._build_url(path) + request_headers = {} if self._session else dict(self.headers) + # HTTP Range is inclusive on both ends, obspec end is exclusive + request_headers["Range"] = f"bytes={start}-{end - 1}" + + async with session.get(url, headers=request_headers) as response: + response.raise_for_status() + return await response.read() + + async def get_range_async( + self, + path: str, + *, + start: int, + end: int | None = None, + length: int | None = None, + ) -> bytes: + """ + Download a byte range asynchronously. + + Parameters + ---------- + path + Path to the file relative to base_url. + start + Start byte offset. + end + End byte offset (exclusive). Either end or length must be provided. + length + Number of bytes to read. Either end or length must be provided. + + Returns + ------- + bytes + The requested byte range. + """ + if end is None and length is None: + raise ValueError("Either 'end' or 'length' must be provided") + if end is None: + end = start + length # type: ignore[operator] + + if self._session is not None: + return await self._do_get_range_async( + self._session, path, start=start, end=end + ) + + # Fallback: create a temporary session for this request + async with aiohttp.ClientSession( + timeout=self.timeout, headers=self.headers + ) as session: + return await self._do_get_range_async(session, path, start=start, end=end) + + async def get_ranges_async( + self, + path: str, + *, + starts: Sequence[int], + ends: Sequence[int] | None = None, + lengths: Sequence[int] | None = None, + ) -> Sequence[bytes]: + """ + Download multiple byte ranges asynchronously. + + Parameters + ---------- + path + Path to the file relative to base_url. + starts + Sequence of start byte offsets. + ends + Sequence of end byte offsets (exclusive). + lengths + Sequence of lengths. Either ends or lengths must be provided. + + Returns + ------- + Sequence[bytes] + The requested byte ranges. + """ + if ends is None and lengths is None: + raise ValueError("Either 'ends' or 'lengths' must be provided") + if ends is None: + ends = [s + ln for s, ln in zip(starts, lengths)] # type: ignore[arg-type] + + if self._session is not None: + # Use managed session for all concurrent requests + tasks = [ + self._do_get_range_async(self._session, path, start=s, end=e) + for s, e in zip(starts, ends) + ] + return await asyncio.gather(*tasks) + + # Fallback: create a single temporary session for all requests + async with aiohttp.ClientSession( + timeout=self.timeout, headers=self.headers + ) as session: + tasks = [ + self._do_get_range_async(session, path, start=s, end=e) + for s, e in zip(starts, ends) + ] + return await asyncio.gather(*tasks) + + # --- Sync methods (wrap async) --- + + def get( + self, + path: str, + *, + options: GetOptions | None = None, + ) -> AiohttpGetResult: + """ + Download a file synchronously. + + This wraps the async implementation for convenience. + + Parameters + ---------- + path + Path to the file relative to base_url. + options + Optional get options. + + Returns + ------- + AiohttpGetResult + Result object with buffer() method and metadata. + """ + result = asyncio.get_event_loop().run_until_complete( + self.get_async(path, options=options) + ) + return AiohttpGetResult( + _data=result._data, + _meta=result._meta, + _attributes=result._attributes, + _range=result._range, + ) + + def get_range( + self, + path: str, + *, + start: int, + end: int | None = None, + length: int | None = None, + ) -> bytes: + """ + Download a byte range synchronously. + + This wraps the async implementation for convenience. + + Parameters + ---------- + path + Path to the file relative to base_url. + start + Start byte offset. + end + End byte offset (exclusive). + length + Number of bytes to read. + + Returns + ------- + bytes + The requested byte range. + """ + return asyncio.get_event_loop().run_until_complete( + self.get_range_async(path, start=start, end=end, length=length) + ) + + def get_ranges( + self, + path: str, + *, + starts: Sequence[int], + ends: Sequence[int] | None = None, + lengths: Sequence[int] | None = None, + ) -> Sequence[bytes]: + """ + Download multiple byte ranges synchronously. + + This wraps the async implementation for convenience. + + Parameters + ---------- + path + Path to the file relative to base_url. + starts + Sequence of start byte offsets. + ends + Sequence of end byte offsets (exclusive). + lengths + Sequence of lengths. + + Returns + ------- + Sequence[bytes] + The requested byte ranges. + """ + return asyncio.get_event_loop().run_until_complete( + self.get_ranges_async(path, starts=starts, ends=ends, lengths=lengths) + ) + + +__all__ = ["AiohttpStore", "AiohttpGetResult", "AiohttpGetResultAsync"] diff --git a/src/obspec_utils/file_handlers.py b/src/obspec_utils/file_handlers.py deleted file mode 100644 index 30c4944..0000000 --- a/src/obspec_utils/file_handlers.py +++ /dev/null @@ -1,75 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING - -import obstore as obs - -if TYPE_CHECKING: - from obstore import ReadableFile - from obstore.store import ObjectStore - -from obstore.store import MemoryStore - - -class ObstoreReader: - _reader: ReadableFile - - def __init__( - self, store: ObjectStore, path: str, buffer_size: int = 1024 * 1024 - ) -> None: - """ - Create an obstore file reader that implements the read, readall, seek, and tell methods, which - can be used in libraries that expect file-like objects. - - This wrapper is necessary in order to return Python bytes types rather than obstore Bytes buffers. - - Parameters - ---------- - store - [ObjectStore][obstore.store.ObjectStore] for reading the file. - path - The path to the file within the store. This should not include the prefix. - buffer_size - The minimum number of bytes to read in a single request. Up to buffer_size bytes will be buffered in memory. - """ - self._reader = obs.open_reader(store, path, buffer_size=buffer_size) - - def read(self, size: int, /) -> bytes: - return self._reader.read(size).to_bytes() - - def readall(self) -> bytes: - return self._reader.read().to_bytes() - - def seek(self, offset: int, whence: int = 0, /): - # TODO: Check on default for whence - return self._reader.seek(offset, whence) - - def tell(self) -> int: - return self._reader.tell() - - -class ObstoreMemCacheReader(ObstoreReader): - _reader: ReadableFile - _memstore: MemoryStore - - def __init__(self, store: ObjectStore, path: str) -> None: - """ - Create an obstore file reader that caches the specified path - in a MemoryStore then performs reads from the file in memory. - - This reader loads the entire file into memory first, which can be beneficial - for files that will be read multiple times or when you want to avoid repeated - network requests to the original store. - - Parameters - ---------- - store - [ObjectStore][obstore.store.ObjectStore] for reading the file. - path - The path to the file within the store. This should not include the prefix. - """ - self._memstore = MemoryStore() - buffer = store.get(path).bytes() - self._memstore.put(path, buffer) - - self._reader = obs.open_reader(self._memstore, path) diff --git a/src/obspec_utils/obspec.py b/src/obspec_utils/obspec.py new file mode 100644 index 0000000..11af5cb --- /dev/null +++ b/src/obspec_utils/obspec.py @@ -0,0 +1,273 @@ +from __future__ import annotations + +import io + + +from typing import Protocol, runtime_checkable + +from obspec import ( + Get, + GetAsync, + GetRange, + GetRangeAsync, + GetRanges, + GetRangesAsync, +) + + +@runtime_checkable +class ReadableStore( + Get, + GetAsync, + GetRange, + GetRangeAsync, + GetRanges, + GetRangesAsync, + Protocol, +): + """ + A minimal protocol for read-only object storage access. + + This protocol defines the intersection of obspec protocols required for + read-only operations like those used by VirtualiZarr. Any object that + implements these methods can be used with ObjectStoreRegistry. + + The protocol includes: + - `get` / `get_async`: Download entire files + - `get_range` / `get_range_async`: Download a single byte range + - `get_ranges` / `get_ranges_async`: Download multiple byte ranges efficiently + + This allows backends like obstore (S3Store, HTTPStore, etc.), aiohttp wrappers, + or any custom implementation to be used interchangeably. + + Examples + -------- + + Using with obstore: + + ```python + from obstore.store import S3Store + from obspec_utils.registry import ObjectStoreRegistry + + # S3Store implements ReadableStore protocol + store = S3Store(bucket="my-bucket") + registry = ObjectStoreRegistry({"s3://my-bucket": store}) + ``` + + Using with a custom aiohttp wrapper: + + ```python + from obspec_utils.registry import ObjectStoreRegistry + from obspec_utils.aiohttp import AiohttpStore + + # AiohttpStore implements ReadableStore protocol + store = AiohttpStore("https://example.com/data") + registry = ObjectStoreRegistry({"https://example.com/data": store}) + ``` + """ + + pass + + +# Protocol-based readers (work with any ReadableStore implementation) + + +class StoreReader: + """ + A file-like reader that works with any ReadableStore protocol implementation. + + This class provides a file-like interface (read, seek, tell) on top of any + object that implements the ReadableStore protocol, including obstore classes, + AiohttpStore, or custom implementations. + + The reader uses `get_range()` calls to fetch data on-demand, with optional + read-ahead buffering for efficiency. + """ + + def __init__( + self, store: ReadableStore, path: str, buffer_size: int = 1024 * 1024 + ) -> None: + """ + Create a file-like reader for any ReadableStore. + + Parameters + ---------- + store + Any object implementing the [ReadableStore][obspec_utils.obspec.ReadableStore] protocol. + path + The path to the file within the store. + buffer_size + Read-ahead buffer size in bytes. When reading, up to this many bytes + may be fetched ahead to reduce the number of requests. + """ + self._store = store + self._path = path + self._buffer_size = buffer_size + self._position = 0 + self._size: int | None = None + # Read-ahead buffer + self._buffer = b"" + self._buffer_start = 0 + + def _get_size(self) -> int: + """Lazily fetch the file size via a get() call.""" + if self._size is None: + result = self._store.get(self._path) + self._size = result.meta["size"] + return self._size + + def read(self, size: int = -1, /) -> bytes: + """ + Read up to `size` bytes from the file. + + Parameters + ---------- + size + Number of bytes to read. If -1, read the entire file. + + Returns + ------- + bytes + The data read from the file. + """ + if size == -1: + return self.readall() + + # Check if we can satisfy from buffer + buffer_end = self._buffer_start + len(self._buffer) + if self._buffer_start <= self._position < buffer_end: + # Some or all data is in buffer + buffer_offset = self._position - self._buffer_start + available = len(self._buffer) - buffer_offset + if available >= size: + # Fully satisfied from buffer + data = self._buffer[buffer_offset : buffer_offset + size] + self._position += len(data) + return data + + # Need to fetch from store + fetch_size = max(size, self._buffer_size) + data = bytes( + self._store.get_range(self._path, start=self._position, length=fetch_size) + ) + + # Update buffer + self._buffer = data + self._buffer_start = self._position + + # Return requested amount + result = data[:size] + self._position += len(result) + return result + + def readall(self) -> bytes: + """ + Read the entire file. + + Returns + ------- + bytes + The complete file contents. + """ + result = self._store.get(self._path) + data = bytes(result.buffer()) + self._size = len(data) + self._position = len(data) + return data + + def seek(self, offset: int, whence: int = 0, /) -> int: + """ + Move the file position. + + Parameters + ---------- + offset + Position offset. + whence + Reference point: 0=start (SEEK_SET), 1=current (SEEK_CUR), 2=end (SEEK_END). + + Returns + ------- + int + The new absolute position. + """ + if whence == 0: # SEEK_SET + self._position = offset + elif whence == 1: # SEEK_CUR + self._position += offset + elif whence == 2: # SEEK_END + self._position = self._get_size() + offset + else: + raise ValueError(f"Invalid whence value: {whence}") + + if self._position < 0: + self._position = 0 + + return self._position + + def tell(self) -> int: + """ + Return the current file position. + + Returns + ------- + int + Current position in bytes from start of file. + """ + return self._position + + +class StoreMemCacheReader: + """ + A file-like reader that caches the entire file in memory. + + This reader fetches the complete file on first access and then serves all + subsequent reads from the in-memory cache. Useful for files that will be + read multiple times or when seeking is frequent. + + Works with any ReadableStore protocol implementation. + """ + + def __init__(self, store: ReadableStore, path: str) -> None: + """ + Create a memory-cached reader for any ReadableStore. + + The file is fetched immediately and cached in memory. + + Parameters + ---------- + store + Any object implementing the [ReadableStore][obspec_utils.obspec.ReadableStore] protocol. + path + The path to the file within the store. + """ + result = store.get(path) + data = bytes(result.buffer()) + self._buffer = io.BytesIO(data) + + def read(self, size: int = -1, /) -> bytes: + """Read up to `size` bytes from the cached file.""" + return self._buffer.read(size) + + def readall(self) -> bytes: + """Read the entire cached file.""" + pos = self._buffer.tell() + self._buffer.seek(0) + data = self._buffer.read() + self._buffer.seek(pos) + return data + + def seek(self, offset: int, whence: int = 0, /) -> int: + """Move the file position within the cached data.""" + return self._buffer.seek(offset, whence) + + def tell(self) -> int: + """Return the current position in the cached file.""" + return self._buffer.tell() + + +__all__: list[str] = [ + "ReadableStore", + "StoreReader", + "StoreMemCacheReader", +] diff --git a/src/obspec_utils/obstore.py b/src/obspec_utils/obstore.py new file mode 100644 index 0000000..3cfda9a --- /dev/null +++ b/src/obspec_utils/obstore.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +try: + import obstore as obs + from obstore.store import MemoryStore +except ImportError as e: + raise ImportError( + "obstore is required for ObstoreReader. Install it with: pip install obstore" + ) from e + + +if TYPE_CHECKING: + from obstore import ReadableFile + from obstore.store import ObjectStore + + +class ObstoreReader: + """ + A file-like reader using obstore's native ReadableFile. + + This class uses obstore's optimized `open_reader()` which provides efficient + buffered reading. It requires an actual [obstore.store.ObjectStore][] instance. + + For a generic reader that works with any ReadableStore, use + [StoreReader][obspec_utils.obspec.StoreReader] instead. + """ + + _reader: ReadableFile + + def __init__( + self, store: ObjectStore, path: str, buffer_size: int = 1024 * 1024 + ) -> None: + """ + Create an obstore file reader. + + Parameters + ---------- + store + An obstore [ObjectStore][obstore.store.ObjectStore] instance. + path + The path to the file within the store. + buffer_size + The minimum number of bytes to read in a single request. + """ + self._reader = obs.open_reader(store, path, buffer_size=buffer_size) + + def read(self, size: int, /) -> bytes: + return self._reader.read(size).to_bytes() + + def readall(self) -> bytes: + return self._reader.read().to_bytes() + + def seek(self, offset: int, whence: int = 0, /) -> int: + return self._reader.seek(offset, whence) + + def tell(self) -> int: + return self._reader.tell() + + +class ObstoreMemCacheReader(ObstoreReader): + """ + A file-like reader using obstore's MemoryStore for caching. + + This class fetches the entire file into obstore's MemoryStore, then uses + obstore's native ReadableFile for efficient cached reads. + + For a generic cached reader that works with any ReadableStore, use + [StoreMemCacheReader][obspec_utils.obspec.StoreMemCacheReader] instead. + """ + + _reader: ReadableFile + _memstore: MemoryStore + + def __init__(self, store: ObjectStore, path: str) -> None: + """ + Create an obstore memory-cached reader. + + Parameters + ---------- + store + An obstore [ObjectStore][obstore.store.ObjectStore] instance. + path + The path to the file within the store. + """ + self._memstore = MemoryStore() + buffer = store.get(path).bytes() + self._memstore.put(path, buffer) + self._reader = obs.open_reader(self._memstore, path) + + +__all__ = ["ObstoreReader", "ObstoreMemCacheReader"] diff --git a/src/obspec_utils/registry.py b/src/obspec_utils/registry.py index 90c2c7a..e690a92 100644 --- a/src/obspec_utils/registry.py +++ b/src/obspec_utils/registry.py @@ -5,16 +5,12 @@ from __future__ import annotations from collections import namedtuple -from typing import TYPE_CHECKING, Dict, Iterator, Optional, Tuple +from typing import Dict, Iterator, Optional, Tuple from urllib.parse import urlparse +from obspec_utils.obspec import ReadableStore from obspec_utils.typing import Path, Url -if TYPE_CHECKING: - from obstore.store import ( - ObjectStore, - ) - UrlKey = namedtuple("UrlKey", ["scheme", "netloc"]) """ A named tuple containing a URL's scheme and authority/netloc. @@ -33,7 +29,7 @@ def get_url_key(url: Url) -> UrlKey: """ Generate the UrlKey containing a url's scheme and authority/netloc that is used a the - primary key's in a [ObjectStoreRegistry.map][obspec_utils.ObjectStoreRegistry.map] + primary key's in a [ObjectStoreRegistry.map][obspec_utils.registry.ObjectStoreRegistry.map] Parameters ---------- @@ -76,10 +72,17 @@ class PathEntry: """ def __init__(self) -> None: - self.store: Optional[ObjectStore] = None + self.store: Optional[ReadableStore] = None self.children: Dict[str, "PathEntry"] = {} - def lookup(self, to_resolve: str) -> Optional[Tuple[ObjectStore, int]]: + def iter_stores(self) -> Iterator[ReadableStore]: + """Iterate over all stores in this entry and its children.""" + if self.store is not None: + yield self.store + for child in self.children.values(): + yield from child.iter_stores() + + def lookup(self, to_resolve: str) -> Optional[Tuple[ReadableStore, int]]: """ Lookup a store based on URL path @@ -103,24 +106,60 @@ def lookup(self, to_resolve: str) -> Optional[Tuple[ObjectStore, int]]: class ObjectStoreRegistry: - def __init__(self, stores: dict[Url, ObjectStore] | None = None) -> None: + """ + A registry that maps URLs to object stores. + + The registry can be used as an async context manager to automatically manage + the lifecycle of stores that support it (like AiohttpStore). Stores that don't + implement the async context manager protocol (like obstore's S3Store) are + unaffected. + + Examples + -------- + + Using as an async context manager with mixed store types: + + ```python + from obstore.store import S3Store + from obspec_utils.registry import ObjectStoreRegistry + from obspec_utils.aiohttp import AiohttpStore + + registry = ObjectStoreRegistry({ + "s3://my-bucket": S3Store(bucket="my-bucket"), + "https://example.com": AiohttpStore("https://example.com"), + }) + + async with registry: + # S3Store works as-is, AiohttpStore session is opened + store, path = registry.resolve("https://example.com/file.nc") + data = await store.get_range_async(path, start=0, end=1000) + # AiohttpStore session is closed automatically + ``` + """ + + def __init__(self, stores: dict[Url, ReadableStore] | None = None) -> None: """ Create a new store registry that matches the provided Urls and - [ObjectStore][obstore.store.ObjectStore] instances. + [ReadableStore][obspec_utils.obspec.ReadableStore] instances. + The registry accepts any object that satisfies the ReadableStore protocol, + which includes obstore classes (S3Store, HTTPStore, etc.) as well as custom + implementations like aiohttp wrappers. Parameters ---------- stores - Mapping of [Url][obspec_utils.typing.Url] to the [ObjectStore][obstore.store.ObjectStore] + Mapping of [Url][obspec_utils.typing.Url] to the [ReadableStore][obspec_utils.obspec.ReadableStore] to be registered under the [Url][obspec_utils.typing.Url]. Examples -------- + Using with obstore: + ```python exec="on" source="above" session="registry-examples" from obstore.store import S3Store - from obspec_utils import ObjectStoreRegistry + from obspec_utils.registry import ObjectStoreRegistry s3store = S3Store(bucket="my-bucket-1", prefix="orig-path") reg = ObjectStoreRegistry({"s3://my-bucket-1": s3store}) @@ -129,6 +168,16 @@ def __init__(self, stores: dict[Url, ObjectStore] | None = None) -> None: assert path == "group/my-file.nc" assert ret is s3store ``` + + Using with any ReadableStore protocol implementation: + + ```python + from obspec_utils.registry import ObjectStoreRegistry + + # Any object implementing the ReadableStore protocol works + custom_store = MyCustomStore("https://example.com/data") + reg = ObjectStoreRegistry({"https://example.com/data": custom_store}) + ``` """ # Mapping from UrlKey (containing scheme and netlocs) to PathEntry self.map: Dict[UrlKey, PathEntry] = {} @@ -136,26 +185,26 @@ def __init__(self, stores: dict[Url, ObjectStore] | None = None) -> None: for url, store in stores.items(): self.register(url, store) - def register(self, url: Url, store: ObjectStore) -> None: + def register(self, url: Url, store: ReadableStore) -> None: """ Register a new store for the provided store [Url][obspec_utils.typing.Url]. - If a store with the same [Url][obspec_utils.typing.Url] existed before, it is replaced. + If a store with the same [Url][obspec_utils.typing.Url] existed before, it is replaced. Parameters ---------- url - [Url][obspec_utils.typing.Url] to registry the [ObjectStore][obstore.store.ObjectStore] under. + [Url][obspec_utils.typing.Url] to register the store under. store - [ObjectStore][obstore.store.ObjectStore] instance to register using the - provided [Url][obspec_utils.typing.Url]. + Any object implementing the [ReadableStore][obspec_utils.obspec.ReadableStore] protocol. + This includes obstore classes (S3Store, HTTPStore, etc.) as well as custom implementations. Examples -------- ```python exec="on" source="above" session="registry-examples" from obstore.store import S3Store - from obspec_utils import ObjectStoreRegistry + from obspec_utils.registry import ObjectStoreRegistry reg = ObjectStoreRegistry() orig_store = S3Store(bucket="my-bucket-1", prefix="orig-path") @@ -182,11 +231,11 @@ def register(self, url: Url, store: ObjectStore) -> None: # Update the store entry.store = store - def resolve(self, url: Url) -> Tuple[ObjectStore, Path]: + def resolve(self, url: Url) -> Tuple[ReadableStore, Path]: """ - Resolve an URL within the [ObjectStoreRegistry][obspec_utils.ObjectStoreRegistry]. + Resolve a URL within the [ObjectStoreRegistry][obspec_utils.registry.ObjectStoreRegistry]. - If [ObjectStoreRegistry.register][obspec_utils.ObjectStoreRegistry.register] has been called + If [ObjectStoreRegistry.register][obspec_utils.registry.ObjectStoreRegistry.register] has been called with a URL with the same scheme and authority/netloc as the object URL, and a path that is a prefix of the provided url's, it is returned along with the trailing path. Paths are matched on a path segment basis, and in the event of multiple possibilities the longest path match is used. @@ -194,20 +243,20 @@ def resolve(self, url: Url) -> Tuple[ObjectStore, Path]: Parameters ---------- url - Url to resolve in the [ObjectStoreRegistry][obspec_utils.ObjectStoreRegistry] + Url to resolve in the [ObjectStoreRegistry][obspec_utils.registry.ObjectStoreRegistry] Returns ------- - ObjectStore - The [ObjectStore][obstore.store.ObjectStore] stored at the resolved url. + ReadableStore + The [ReadableStore][obspec_utils.obspec.ReadableStore] stored at the resolved url. Path The trailing portion of the url after the prefix of the matching store in the - [ObjectStoreRegistry][obspec_utils.ObjectStoreRegistry]. + [ObjectStoreRegistry][obspec_utils.registry.ObjectStoreRegistry]. Raises ------ ValueError - If the URL cannot be resolved, meaning that [ObjectStoreRegistry.register][obspec_utils.ObjectStoreRegistry.register] + If the URL cannot be resolved, meaning that [ObjectStoreRegistry.register][obspec_utils.registry.ObjectStoreRegistry.register] has not been called with a URL with the same scheme and authority/netloc as the object URL, and a path that is a prefix of the provided url's. @@ -216,7 +265,7 @@ def resolve(self, url: Url) -> Tuple[ObjectStore, Path]: ```python exec="on" source="above" session="registry-resolve-examples" from obstore.store import MemoryStore, S3Store - from obspec_utils import ObjectStoreRegistry + from obspec_utils.registry import ObjectStoreRegistry registry = ObjectStoreRegistry() memstore1 = MemoryStore() @@ -272,6 +321,55 @@ def resolve(self, url: Url) -> Tuple[ObjectStore, Path]: return store, path_after_prefix raise ValueError(f"Could not find an ObjectStore matching the url `{url}`") + def _iter_stores(self) -> Iterator[ReadableStore]: + """Iterate over all registered stores.""" + for entry in self.map.values(): + yield from entry.iter_stores() + + async def __aenter__(self) -> "ObjectStoreRegistry": + """ + Enter the async context manager, opening all stores that support it. + + Stores that implement the async context manager protocol (like AiohttpStore) + will have their sessions initialized. Stores that don't support it (like + obstore's S3Store) are unaffected. + + Examples + -------- + + ```python + from obstore.store import S3Store + from obspec_utils.registry import ObjectStoreRegistry + from obspec_utils.aiohttp import AiohttpStore + + registry = ObjectStoreRegistry({ + "s3://my-bucket": S3Store(bucket="my-bucket"), + "https://example.com": AiohttpStore("https://example.com"), + }) + + async with registry: + # S3Store works as-is, AiohttpStore session is opened + store, path = registry.resolve("https://example.com/file.nc") + data = await store.get_range_async(path, start=0, end=1000) + # AiohttpStore session is closed + ``` + """ + for store in self._iter_stores(): + if hasattr(store, "__aenter__"): + await store.__aenter__() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: + """ + Exit the async context manager, closing all stores that support it. + + Stores that implement the async context manager protocol will have their + resources cleaned up. Stores that don't support it are unaffected. + """ + for store in self._iter_stores(): + if hasattr(store, "__aexit__"): + await store.__aexit__(exc_type, exc_val, exc_tb) + def path_segments(path: str) -> Iterator[str]: """ diff --git a/tests/test_registry.py b/tests/test_registry.py index 8a54276..8d7c936 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -1,7 +1,8 @@ import pytest from obstore.store import MemoryStore -from obspec_utils import ObjectStoreRegistry +from obspec_utils.registry import ObjectStoreRegistry +from obspec_utils.obspec import ReadableStore def test_registry(): @@ -32,3 +33,223 @@ def test_resolve_raises(): ): url = "s3://bucket1/path/to/object" ret, path = registry.resolve(url) + + +def test_obstore_satisfies_readable_store_protocol(): + """Verify that obstore classes satisfy the ReadableStore protocol.""" + memstore = MemoryStore() + # Runtime check using isinstance with @runtime_checkable protocol + assert isinstance(memstore, ReadableStore) + + +def test_registry_with_custom_readable_store(): + """Test that registry works with any ReadableStore protocol implementation.""" + from collections.abc import Sequence + + class MockReadableStore: + """A minimal mock that satisfies the ReadableStore protocol.""" + + def __init__(self, data: bytes = b"test data"): + self._data = data + + def get(self, path, *, options=None): + # Return a mock GetResult + return MockGetResult(self._data) + + async def get_async(self, path, *, options=None): + return MockGetResultAsync(self._data) + + def get_range(self, path, *, start, end=None, length=None): + if end is None: + end = start + length + return self._data[start:end] + + async def get_range_async(self, path, *, start, end=None, length=None): + if end is None: + end = start + length + return self._data[start:end] + + def get_ranges( + self, path, *, starts, ends=None, lengths=None + ) -> Sequence[bytes]: + if ends is None: + ends = [s + ln for s, ln in zip(starts, lengths)] + return [self._data[s:e] for s, e in zip(starts, ends)] + + async def get_ranges_async( + self, path, *, starts, ends=None, lengths=None + ) -> Sequence[bytes]: + if ends is None: + ends = [s + ln for s, ln in zip(starts, lengths)] + return [self._data[s:e] for s, e in zip(starts, ends)] + + class MockGetResult: + def __init__(self, data): + self._data = data + + @property + def attributes(self): + return {} + + def buffer(self): + return self._data + + @property + def meta(self): + return { + "path": "", + "last_modified": None, + "size": len(self._data), + "e_tag": None, + "version": None, + } + + @property + def range(self): + return (0, len(self._data)) + + def __iter__(self): + yield self._data + + class MockGetResultAsync: + def __init__(self, data): + self._data = data + + @property + def attributes(self): + return {} + + async def buffer_async(self): + return self._data + + @property + def meta(self): + return { + "path": "", + "last_modified": None, + "size": len(self._data), + "e_tag": None, + "version": None, + } + + @property + def range(self): + return (0, len(self._data)) + + async def __aiter__(self): + yield self._data + + # Create a mock store and register it + mock_store = MockReadableStore(b"hello world") + registry = ObjectStoreRegistry({"https://example.com": mock_store}) + + # Resolve and use the store + store, path = registry.resolve("https://example.com/data/file.txt") + assert store is mock_store + assert path == "data/file.txt" + + # Verify the protocol methods work + assert store.get_range(path, start=0, end=5) == b"hello" + assert store.get_range(path, start=6, length=5) == b"world" + + +@pytest.mark.asyncio +async def test_registry_with_async_operations(): + """Test async operations with registry.""" + memstore = MemoryStore() + # Put some test data + memstore.put("test.txt", b"async test data") + + registry = ObjectStoreRegistry({"mem://test": memstore}) + store, path = registry.resolve("mem://test/test.txt") + + # Test async get_range + result = await store.get_range_async(path, start=0, end=5) + assert bytes(result) == b"async" + + +class TestStoreReader: + """Tests for the generic StoreReader class.""" + + def test_store_reader_with_obstore(self): + """Test StoreReader with an obstore MemoryStore.""" + from obspec_utils.obspec import StoreReader + + memstore = MemoryStore() + memstore.put("test.txt", b"hello world from store reader") + + reader = StoreReader(memstore, "test.txt", buffer_size=10) + + # Test read + assert reader.read(5) == b"hello" + assert reader.tell() == 5 + + # Test seek and read + reader.seek(6) + assert reader.read(5) == b"world" + + # Test seek from current + reader.seek(-5, 1) # SEEK_CUR + assert reader.read(5) == b"world" + + # Test readall + reader.seek(0) + assert reader.readall() == b"hello world from store reader" + + def test_store_reader_seek_end(self): + """Test SEEK_END functionality.""" + from obspec_utils.obspec import StoreReader + + memstore = MemoryStore() + memstore.put("test.txt", b"0123456789") + + reader = StoreReader(memstore, "test.txt") + + # Seek to 2 bytes before end + reader.seek(-2, 2) # SEEK_END + assert reader.read(2) == b"89" + + def test_store_reader_buffering(self): + """Test that buffering works correctly.""" + from obspec_utils.obspec import StoreReader + + memstore = MemoryStore() + memstore.put("test.txt", b"0123456789ABCDEF") + + # Small buffer size to test buffering behavior + reader = StoreReader(memstore, "test.txt", buffer_size=8) + + # First read should fetch buffer_size bytes + assert reader.read(2) == b"01" + # Second read should come from buffer + assert reader.read(2) == b"23" + + +class TestStoreMemCacheReader: + """Tests for the generic StoreMemCacheReader class.""" + + def test_store_memcache_reader(self): + """Test StoreMemCacheReader with an obstore MemoryStore.""" + from obspec_utils.obspec import StoreMemCacheReader + + memstore = MemoryStore() + memstore.put("test.txt", b"cached content here") + + reader = StoreMemCacheReader(memstore, "test.txt") + + # Test read + assert reader.read(6) == b"cached" + + # Test seek and read + reader.seek(7) + assert reader.read(7) == b"content" + + # Test readall preserves position + pos_before = reader.tell() + data = reader.readall() + assert data == b"cached content here" + assert reader.tell() == pos_before + + # Test seek to start + reader.seek(0) + assert reader.read() == b"cached content here" diff --git a/tests/test_xarray.py b/tests/test_xarray.py index 612deaa..0014f12 100644 --- a/tests/test_xarray.py +++ b/tests/test_xarray.py @@ -1,5 +1,5 @@ import xarray as xr -from obspec_utils import ObstoreMemCacheReader, ObstoreReader +from obspec_utils.obstore import ObstoreMemCacheReader, ObstoreReader from obstore.store import LocalStore