From 3fe675b8ae7bd0efd6bcb7c1ad815592aab57164 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Tue, 16 Jun 2026 23:17:19 -0400 Subject: [PATCH 01/10] wip: python-implementable store protocols --- src/error.rs | 8 ++++++-- src/storage/key.rs | 40 ++++++++++++++++++++++++++++++++++++++++ src/storage/mod.rs | 2 ++ src/storage/python.rs | 17 +++++++++++++++++ 4 files changed, 65 insertions(+), 2 deletions(-) create mode 100644 src/storage/key.rs create mode 100644 src/storage/python.rs diff --git a/src/error.rs b/src/error.rs index f204c1d..efd8338 100644 --- a/src/error.rs +++ b/src/error.rs @@ -15,7 +15,7 @@ use zarrs::array::{ArrayCreateError, ArrayError}; use zarrs::filesystem::FilesystemStoreCreateError; use zarrs::group::GroupCreateError; use zarrs::node::{NodeCreateError, NodePathError}; -use zarrs::storage::StorageError; +use zarrs::storage::{StorageError, StoreKeyError}; create_exception!( zarrista, @@ -38,7 +38,7 @@ create_exception!( /// appropriate Python exception. #[derive(Debug, Error)] #[non_exhaustive] -pub(crate) enum ZarristaError { +pub enum ZarristaError { /// No array or group exists at the requested path. #[error("{0}")] NotFound(String), @@ -72,6 +72,9 @@ pub(crate) enum ZarristaError { /// Failed to (de)serialize JSON. #[error(transparent)] SerdeJson(#[from] serde_json::Error), + /// Store key error + #[error(transparent)] + StoreKey(#[from] StoreKeyError), } impl ZarristaError { @@ -97,6 +100,7 @@ impl From for PyErr { ZarristaException::new_err(err.to_string()) } ZarristaError::SerdeJson(err) => ZarristaException::new_err(err.to_string()), + ZarristaError::StoreKey(err) => ZarristaException::new_err(err.to_string()), } } } diff --git a/src/storage/key.rs b/src/storage/key.rs new file mode 100644 index 0000000..c505953 --- /dev/null +++ b/src/storage/key.rs @@ -0,0 +1,40 @@ +use std::convert::Infallible; + +use pyo3::prelude::*; +use pyo3::types::PyString; +use zarrs::storage::StoreKey; + +use crate::error::ZarristaError; + +pub struct PyStoreKey(StoreKey); + +impl FromPyObject<'_, '_> for PyStoreKey { + type Error = ZarristaError; + + fn extract(obj: Borrowed<'_, '_, PyAny>) -> Result { + let key = obj.extract::()?; + Ok(PyStoreKey(StoreKey::new(key)?)) + } +} + +impl<'py> IntoPyObject<'py> for PyStoreKey { + type Target = PyString; + type Error = Infallible; + type Output = Bound<'py, Self::Target>; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(PyString::new(py, self.0.as_str())) + } +} + +impl From for StoreKey { + fn from(key: PyStoreKey) -> Self { + key.0 + } +} + +impl From for PyStoreKey { + fn from(key: StoreKey) -> Self { + PyStoreKey(key) + } +} diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 2f3f993..f14e743 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -1,4 +1,6 @@ mod r#async; +mod key; +mod python; mod sync; #[allow(unused)] diff --git a/src/storage/python.rs b/src/storage/python.rs new file mode 100644 index 0000000..b953530 --- /dev/null +++ b/src/storage/python.rs @@ -0,0 +1,17 @@ +// use pyo3::prelude::*; +// use zarrs::storage::{ReadableStorage, ReadableStorageTraits, StorageError, StoreKey}; + +// use crate::storage::key::PyStoreKey; + +// /// A Python backend for making requests that conforms to the GetRangeAsync and GetRangesAsync +// /// protocols defined by obspec. +// /// https://developmentseed.org/obspec/latest/api/get/#obspec.GetRangeAsync +// /// https://developmentseed.org/obspec/latest/api/get/#obspec.GetRangesAsync +// #[derive(Debug)] +// pub(crate) struct PyReadable(Py); + +// impl ReadableStorageTraits for PyReadable { +// fn size_key(&self, key: &StoreKey) -> Result, StorageError> { +// let key = PyStoreKey::from(key.clone()); +// } +// } From 9429c613f05a29a49a9e10450bbc17f64efd578e Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Tue, 16 Jun 2026 23:35:48 -0400 Subject: [PATCH 02/10] Add design spec for custom Python store protocol Sync, duck-typed Python stores adapted to zarrs ReadableListableStorageTraits via a single PyStore wrapper that degrades listing at runtime. Capabilities declared via @property predicates (supports_get_partial / supports_listing). Co-Authored-By: Claude Opus 4.8 (1M context) --- ...2026-06-16-python-store-protocol-design.md | 187 ++++++++++++++++++ 1 file changed, 187 insertions(+) create mode 100644 dev-docs/specs/2026-06-16-python-store-protocol-design.md diff --git a/dev-docs/specs/2026-06-16-python-store-protocol-design.md b/dev-docs/specs/2026-06-16-python-store-protocol-design.md new file mode 100644 index 0000000..5f2d1e9 --- /dev/null +++ b/dev-docs/specs/2026-06-16-python-store-protocol-design.md @@ -0,0 +1,187 @@ +# Custom Python store protocol (sync) + +**Date:** 2026-06-16 +**Status:** Approved — ready for implementation plan + +## Goal + +Let a user pass an arbitrary, duck-typed Python object as a store, instead of +only the built-in `FilesystemStore` / `MemoryStore`. The object declares its +capabilities (partial reads, listing) and implements a small set of methods; +zarrista adapts it to the `zarrs` storage traits so it works anywhere the +built-in stores do. + +This pulls forward roadmap Tier-3 item 7b ("a Python-implementable `Store` +protocol for custom backends, mirroring zarr-python's `Store` ABC"). See +[2026-06-16-api-roadmap.md](2026-06-16-api-roadmap.md). + +### In scope + +- **Sync only.** Implement `zarrs` `ReadableStorageTraits` + + `ListableStorageTraits` over a Python object. +- **Readable required, listable optional**, detected via capability predicates. + +### Out of scope (deferred, not foreclosed) + +- **Async** custom stores (awaiting Python coroutines via + `pyo3-async-runtimes`). The async/obstore path already exists and remains + where real I/O concurrency lives. +- **Writable / deletable** custom stores. Writing is out of scope project-wide + today; when it lands it forces a different `zarrs` static type + (`ReadableWritableStorageTraits`) and is the right moment to introduce a + storage enum or split path. We do **not** build that now. + +## Why mirror `zarrs`, with an obspec adapter in Python + +The Rust↔Python boundary mirrors the methods `zarrs` actually needs, smoothed +into Pythonic shapes. The compiled core stays decoupled from any external +spec; an obspec/obstore adapter can live in pure Python (iterate without +recompiling), and obstore users still get a one-liner. Method names mirror +`zarrs` (which mostly coincides with zarr-python's `Store` ABC); the one +borrowed-from-zarr-python concept is the `supports_listing` predicate, because +`zarrs` expresses listability through its *type system* rather than a method — +something we cannot replicate dynamically. + +## Architecture + +The entire codebase is statically typed on a single trait object, +`Arc`. `Array::open` and `Group::open` both +take exactly that, and `Group::array_keys()` / `group_keys()` go through +`child_array_paths()`, which requires listing. + +Rather than refactor `Array`/`Group` onto a capability enum (which fights +`zarrs`'s static generics and forces parallel `Array`/`Group` instantiations +per variant), use **a single Rust wrapper that always implements +`ReadableListableStorageTraits` and degrades at runtime**: + +```rust +// src/storage/python.rs +pub(crate) struct PyStore(Py); // Py is Send + Sync + +impl ReadableStorageTraits for PyStore { /* calls Python get / get_partial* / size_key */ } + +impl ListableStorageTraits for PyStore { + fn list(&self) -> Result { + // supports_listing == false -> Err(StorageError::Unsupported(...)) + // else call the Python list method + } + // list_prefix / list_dir / size_prefix likewise +} +// blanket impl => ReadableListableStorageTraits for free +``` + +`extract_storage` grows one arm: anything that isn't a `FilesystemStore` / +`MemoryStore` is wrapped as `PyStore` and returned as +`Arc`. **No changes to `Array`, `Group`, +`node`, or any call site.** A readable-only Python store opens arrays fine; +`group.array_keys()` on it raises a clear runtime error. + +This replaces the unused `SyncStorage` enum in `src/storage/sync.rs` for now; +the enum returns when writing is implemented and genuinely needs it. + +## The Python protocol (sync) + +A duck-typed object. Capabilities are declared via `@property` (zarr-python +style); methods provide the bytes/keys. Tiered so a trivial store (e.g. a dict) +needs almost nothing, while a real backend opts into efficient partial reads. + +### Capability predicates (`@property`, authoritative) + +| Property | Type | Meaning | +| ------------------------ | ------ | ---------------------------------------------------- | +| `supports_get_partial` | `bool` | drives `zarrs` `supports_get_partial()`; gates the partial-read methods | +| `supports_listing` | `bool` | gates the listable methods | + +A missing property is treated as `False`. + +### Readable + +| Method | Required? | Returns | +| --------------------------------------------------------- | ------------------------------- | ----------------------------- | +| `get(key: str)` | **yes** | `bytes \| None` (None = absent) | +| `get_partial(key: str, byte_range: ByteRange)` | optional (`supports_get_partial`) | `bytes \| None` | +| `get_partial_many(key: str, byte_ranges: list[ByteRange])`| optional (`supports_get_partial`) | `list[bytes] \| None` | +| `size_key(key: str)` | optional | `int \| None` | + +- `get` alone yields a correct, working store. +- When `supports_get_partial` is `False`, Rust synthesizes partial reads by + fetching the full value via `get` and slicing. Efficiency is opt-in; + correctness is free. +- `get_partial_many` is the one method `zarrs` strictly requires + (`get_partial_many` on the trait). If the Python object provides only + `get_partial`, Rust loops it; if it provides neither, Rust falls back to + `get` + slice. +- `size_key` absent → Rust falls back to `len(get(key))`. + +### Listable (only consulted when `supports_listing`) + +| Method | Returns | +| --------------------------- | --------------------------------------------- | +| `list()` | `list[str]` | +| `list_prefix(prefix: str)` | `list[str]` | +| `list_dir(prefix: str)` | `{"keys": list[str], "prefixes": list[str]}` | +| `size_prefix(prefix: str)` | `int` | + +`list_dir` mirrors `zarrs`'s `list_dir`, which returns both the keys and the +child prefixes directly under `prefix`. + +### `ByteRange` representation + +`zarrs`'s `ByteRange` is an enum: `FromStart(offset, Option)` **or** +`FromEnd(offset, Option)` (suffix reads, which sharding uses heavily). +The Python representation must express both ends — a start-only tuple would +make suffix reads inexpressible. + +Representation: a small object/tuple carrying `(anchor, offset, length)` where +`anchor` is `"start"` or `"end"` and `length` may be `None` (to end of value). +The exact concrete form (lightweight dataclass vs. tuple vs. exported pyclass) +is an implementation choice for the plan; the constraint is that both anchors +and an optional length round-trip. + +## Error handling + +Mirror zarr-python's categories as closely as the bridge allows: + +- **Missing key** → the Python method returns `None` → `zarrs` `None` + (not-found). No exception. Matches zarr-python `get` semantics. +- **Unsupported capability** (e.g. listing when `supports_listing` is `False`) + → `StorageError::Unsupported("store does not support listing")`, short- + circuited before any Python call. Surfaced to Python as a clear, + zarr-python-like "operation not supported" error. +- **Exception raised inside a store method** → caught and wrapped as + `StorageError::Other()`, preserving the message. + **Known limitation:** `zarrs`'s `StorageError` carries only strings, so the + original Python exception *type* and traceback are flattened to a message on + the way back out. Acceptable for this milestone; revisit only if it bites. + +## GIL / threading + +`Py` is `Send + Sync`, so `PyStore` satisfies the trait bounds. Every +method enters via `Python::attach` to take the GIL. Consequence, stated +plainly and accepted: `zarrs` may call a sync store from multiple rayon +threads, but the GIL serializes the Python-side calls — a sync custom Python +store does **not** get true I/O parallelism. That is inherent to sync-first; +the async/obstore path remains where concurrency lives. + +## Testing + +1. A pure-Python dict-backed store implementing the protocol, used to open an + array/group and read a chunk — proves the boundary end-to-end. +2. A readable-only store (no `supports_listing`) asserting `group.array_keys()` + raises the clear "not supported" error. +3. A Rust unit test under `auto-initialize` defining a tiny store class inline. + +Round-trip tests vs. zarr-python can come with the broader harness the roadmap +already calls for. + +## Files touched + +- `src/storage/python.rs` — new `PyStore` wrapper + trait impls (replaces the + current commented stub). +- `src/storage/sync.rs` — `extract_storage` grows the `PyStore` arm; remove the + unused `SyncStorage` enum (returns with writing). +- `src/storage/mod.rs` — export wiring as needed. +- `src/error.rs` — map `StorageError::Unsupported` to a clear Python exception + if the default `ZarristaException` message isn't sufficient. +- Tests under `tests/` (Python) and an inline Rust unit test. + From a3ae2f4c2a1cf699d4e8ad8678403b86606bcc5c Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Tue, 16 Jun 2026 23:47:24 -0400 Subject: [PATCH 03/10] Fix ByteRange variants in store protocol spec ByteRange is FromStart/Suffix, not FromStart/FromEnd. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../2026-06-16-python-store-protocol-design.md | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/dev-docs/specs/2026-06-16-python-store-protocol-design.md b/dev-docs/specs/2026-06-16-python-store-protocol-design.md index 5f2d1e9..95b24cb 100644 --- a/dev-docs/specs/2026-06-16-python-store-protocol-design.md +++ b/dev-docs/specs/2026-06-16-python-store-protocol-design.md @@ -127,13 +127,15 @@ child prefixes directly under `prefix`. ### `ByteRange` representation -`zarrs`'s `ByteRange` is an enum: `FromStart(offset, Option)` **or** -`FromEnd(offset, Option)` (suffix reads, which sharding uses heavily). -The Python representation must express both ends — a start-only tuple would -make suffix reads inexpressible. - -Representation: a small object/tuple carrying `(anchor, offset, length)` where -`anchor` is `"start"` or `"end"` and `length` may be `None` (to end of value). +`zarrs`'s `ByteRange` is an enum: `FromStart(offset: u64, Option)` +**or** `Suffix(length: u64)` (a read of the last `length` bytes, which sharding +uses heavily). The Python representation must express both — a start-only tuple +would make suffix reads inexpressible. + +Representation: a small object/tuple carrying `(kind, offset, length)` where +`kind` is `"start"` or `"suffix"` and, for `"start"`, `length` may be `None` +(read to the end of the value). For `"suffix"`, `offset` is unused and `length` +is the suffix size. The exact concrete form (lightweight dataclass vs. tuple vs. exported pyclass) is an implementation choice for the plan; the constraint is that both anchors and an optional length round-trip. From a0e524b05b53f6683694cd4cd95eb4162fb1e97f Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Tue, 16 Jun 2026 23:54:04 -0400 Subject: [PATCH 04/10] feat: readable + listable custom Python store via PyStore Wrap an arbitrary duck-typed Python object as a sync zarrs store. extract_storage now falls through to PyStore for any non-builtin object. Listing degrades to a clear "not supported" error when not declared. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/storage/mod.rs | 3 +- src/storage/python.rs | 227 ++++++++++++++++++++++++++++++++++--- src/storage/sync.rs | 16 +-- tests/test_custom_store.py | 46 ++++++++ 4 files changed, 261 insertions(+), 31 deletions(-) create mode 100644 tests/test_custom_store.py diff --git a/src/storage/mod.rs b/src/storage/mod.rs index f14e743..ca4aef4 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -3,8 +3,9 @@ mod key; mod python; mod sync; +pub(crate) use python::PyStore; #[allow(unused)] pub use r#async::AsyncStorage; pub(crate) use sync::extract_storage; #[allow(unused)] -pub use sync::{PyFilesystemStore, PyMemoryStore, SyncStorage}; +pub use sync::{PyFilesystemStore, PyMemoryStore}; diff --git a/src/storage/python.rs b/src/storage/python.rs index b953530..3ce9b99 100644 --- a/src/storage/python.rs +++ b/src/storage/python.rs @@ -1,17 +1,210 @@ -// use pyo3::prelude::*; -// use zarrs::storage::{ReadableStorage, ReadableStorageTraits, StorageError, StoreKey}; - -// use crate::storage::key::PyStoreKey; - -// /// A Python backend for making requests that conforms to the GetRangeAsync and GetRangesAsync -// /// protocols defined by obspec. -// /// https://developmentseed.org/obspec/latest/api/get/#obspec.GetRangeAsync -// /// https://developmentseed.org/obspec/latest/api/get/#obspec.GetRangesAsync -// #[derive(Debug)] -// pub(crate) struct PyReadable(Py); - -// impl ReadableStorageTraits for PyReadable { -// fn size_key(&self, key: &StoreKey) -> Result, StorageError> { -// let key = PyStoreKey::from(key.clone()); -// } -// } +//! A custom, duck-typed Python object adapted to the `zarrs` sync storage traits. +//! +//! The Python object declares capabilities via `@property` predicates and +//! implements a small set of methods. [`PyStore`] reads the capability flags +//! once at construction and adapts the object to [`ReadableStorageTraits`] +//! (+ [`ListableStorageTraits`], added in a later task). + +use pyo3::call::PyCallArgs; +use pyo3::prelude::*; +use zarrs::storage::byte_range::{ByteRange, ByteRangeIterator}; +use zarrs::storage::{ + Bytes, ListableStorageTraits, MaybeBytes, MaybeBytesIterator, ReadableStorageTraits, + StorageError, StoreKey, StoreKeys, StoreKeysPrefixes, StorePrefix, StorePrefixes, +}; + +/// A Python object adapted to the `zarrs` sync storage traits. +#[derive(Debug)] +pub(crate) struct PyStore { + obj: Py, + supports_get_partial: bool, + supports_listing: bool, +} + +impl PyStore { + /// Wrap a duck-typed Python store object, reading its capability flags now. + pub(crate) fn new(obj: &Bound<'_, PyAny>) -> Self { + Self { + obj: obj.clone().unbind(), + supports_get_partial: read_bool_property(obj, "supports_get_partial"), + supports_listing: read_bool_property(obj, "supports_listing"), + } + } + + /// Call the Python `get(key)` method, returning the full value or `None`. + fn py_get(&self, key: &StoreKey) -> Result { + Python::attach(|py| { + let result = self + .obj + .bind(py) + .call_method1("get", (key.as_str(),)) + .map_err(py_to_storage_error)?; + if result.is_none() { + return Ok(None); + } + let bytes = result.extract::>().map_err(py_to_storage_error)?; + Ok(Some(Bytes::from(bytes))) + }) + } +} + +/// Read a boolean `@property`; a missing or non-bool property is `false`. +fn read_bool_property(obj: &Bound<'_, PyAny>, name: &str) -> bool { + obj.getattr(name) + .ok() + .and_then(|v| v.extract::().ok()) + .unwrap_or(false) +} + +/// Map a Python exception to a `zarrs` `StorageError`. +/// +/// `zarrs`'s `StorageError` carries only strings, so the Python exception type +/// and traceback are flattened to the exception's display string here. +fn py_to_storage_error(err: PyErr) -> StorageError { + StorageError::Other(err.to_string()) +} + +impl ReadableStorageTraits for PyStore { + fn get(&self, key: &StoreKey) -> Result { + self.py_get(key) + } + + fn get_partial_many<'a>( + &'a self, + key: &StoreKey, + byte_ranges: ByteRangeIterator<'a>, + ) -> Result, StorageError> { + // MVP: fetch the full value and slice each range (mirrors MemoryStore). + let Some(data) = self.py_get(key)? else { + return Ok(None); + }; + let out = byte_ranges.map(move |byte_range: ByteRange| { + let len = data.len() as u64; + let start = byte_range.start(len) as usize; + let end = byte_range.end(len) as usize; + if end > data.len() { + Err(StorageError::Other(format!( + "byte range {byte_range:?} out of bounds for value of length {}", + data.len() + ))) + } else { + Ok(data.slice(start..end)) + } + }); + Ok(Some(Box::new(out))) + } + + fn size_key(&self, key: &StoreKey) -> Result, StorageError> { + // Prefer an explicit `size_key`; else fall back to len(get(key)). + let has_size_key = Python::attach(|py| { + self.obj + .bind(py) + .hasattr("size_key") + .map_err(py_to_storage_error) + })?; + if has_size_key { + return Python::attach(|py| { + let result = self + .obj + .bind(py) + .call_method1("size_key", (key.as_str(),)) + .map_err(py_to_storage_error)?; + if result.is_none() { + return Ok(None); + } + result + .extract::() + .map(Some) + .map_err(py_to_storage_error) + }); + } + Ok(self.py_get(key)?.map(|b| b.len() as u64)) + } + + fn supports_get_partial(&self) -> bool { + self.supports_get_partial + } +} + +impl PyStore { + /// Error returned by every listable method when listing is not declared. + fn require_listing(&self) -> Result<(), StorageError> { + if self.supports_listing { + Ok(()) + } else { + Err(StorageError::Unsupported( + "store does not support listing".to_string(), + )) + } + } + + /// Call a Python method returning a list of key strings. + fn py_list_keys<'py, A: PyCallArgs<'py>>( + &self, + py: Python<'py>, + method: &str, + args: A, + ) -> Result { + let result = self + .obj + .bind(py) + .call_method1(method, args) + .map_err(py_to_storage_error)?; + let raw = result + .extract::>() + .map_err(py_to_storage_error)?; + raw.into_iter() + .map(|k| StoreKey::new(k).map_err(|e| StorageError::Other(e.to_string()))) + .collect() + } +} + +impl ListableStorageTraits for PyStore { + fn list(&self) -> Result { + self.require_listing()?; + Python::attach(|py| self.py_list_keys(py, "list", ())) + } + + fn list_prefix(&self, prefix: &StorePrefix) -> Result { + self.require_listing()?; + Python::attach(|py| self.py_list_keys(py, "list_prefix", (prefix.as_str(),))) + } + + fn list_dir(&self, prefix: &StorePrefix) -> Result { + self.require_listing()?; + Python::attach(|py| { + let result = self + .obj + .bind(py) + .call_method1("list_dir", (prefix.as_str(),)) + .map_err(py_to_storage_error)?; + let keys_obj = result.get_item("keys").map_err(py_to_storage_error)?; + let prefixes_obj = result.get_item("prefixes").map_err(py_to_storage_error)?; + let keys = keys_obj + .extract::>() + .map_err(py_to_storage_error)? + .into_iter() + .map(|k| StoreKey::new(k).map_err(|e| StorageError::Other(e.to_string()))) + .collect::>()?; + let prefixes = prefixes_obj + .extract::>() + .map_err(py_to_storage_error)? + .into_iter() + .map(|p| StorePrefix::new(p).map_err(|e| StorageError::Other(e.to_string()))) + .collect::>()?; + Ok(StoreKeysPrefixes::new(keys, prefixes)) + }) + } + + fn size_prefix(&self, prefix: &StorePrefix) -> Result { + self.require_listing()?; + Python::attach(|py| { + self.obj + .bind(py) + .call_method1("size_prefix", (prefix.as_str(),)) + .map_err(py_to_storage_error)? + .extract::() + .map_err(py_to_storage_error) + }) + } +} diff --git a/src/storage/sync.rs b/src/storage/sync.rs index 6e4ccf3..0a1404d 100644 --- a/src/storage/sync.rs +++ b/src/storage/sync.rs @@ -1,5 +1,5 @@ use crate::error::ZarristaResult; -use pyo3::exceptions::PyTypeError; +use crate::storage::PyStore; use pyo3::prelude::*; use std::path::PathBuf; use std::sync::Arc; @@ -7,15 +7,6 @@ use zarrs::filesystem::FilesystemStore; use zarrs::storage::store::MemoryStore; use zarrs::storage::ReadableListableStorageTraits; -use zarrs::storage::{ReadableStorageTraits, ReadableWritableStorageTraits, WritableStorageTraits}; - -#[allow(dead_code)] -pub enum SyncStorage { - Readable(Arc), - Writable(Arc), - ReadableWritable(Arc), -} - /// A store backed by a local directory. #[pyclass(module = "zarrista", frozen, name = "FilesystemStore")] pub struct PyFilesystemStore { @@ -68,7 +59,6 @@ pub(crate) fn extract_storage( if let Ok(s) = store.cast::() { return Ok(s.get().storage.clone()); } - Err(PyTypeError::new_err( - "expected a FilesystemStore or MemoryStore", - )) + // Any other object is treated as a duck-typed custom store. + Ok(Arc::new(PyStore::new(store))) } diff --git a/tests/test_custom_store.py b/tests/test_custom_store.py new file mode 100644 index 0000000..be2354e --- /dev/null +++ b/tests/test_custom_store.py @@ -0,0 +1,46 @@ +import os + +import numpy as np +import pytest +import zarr + +from zarrista import Array + + +@pytest.fixture +def zarr_bytes(tmp_path): + """Write a tiny v3 array with zarr-python, return {key: bytes}.""" + store = zarr.storage.LocalStore(str(tmp_path)) + root = zarr.create_group(store=store) + arr = root.create_array("a", shape=(4,), chunks=(2,), dtype="int32") + arr[:] = np.arange(4, dtype="int32") + + mapping: dict[str, bytes] = {} + for dirpath, _dirs, files in os.walk(tmp_path): + for name in files: + full = os.path.join(dirpath, name) + rel = os.path.relpath(full, tmp_path).replace(os.sep, "/") + with open(full, "rb") as fh: + mapping[rel] = fh.read() + return mapping + + +class ReadOnlyDictStore: + """Minimal readable store: implements only `get`.""" + + supports_get_partial = False + supports_listing = False + + def __init__(self, mapping: dict[str, bytes]): + self._mapping = mapping + + def get(self, key: str) -> bytes | None: + return self._mapping.get(key) + + +def test_open_array_and_read_chunk_from_custom_store(zarr_bytes): + store = ReadOnlyDictStore(zarr_bytes) + array = Array.open(store, "/a") + data = array.retrieve_chunk([0]) + assert array.shape == [4] + np.testing.assert_array_equal(data.to_numpy(), np.array([0, 1], dtype="int32")) From cb9946bb47c48f4fe4010ee4680b45f522c0c3f0 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Wed, 17 Jun 2026 00:11:44 -0400 Subject: [PATCH 05/10] Create PySyncStorage type --- src/array/sync.rs | 7 +++---- src/group/sync.rs | 10 +++++----- src/storage/mod.rs | 4 +--- src/storage/sync.rs | 38 ++++++++++++++++++++++++-------------- 4 files changed, 33 insertions(+), 26 deletions(-) diff --git a/src/array/sync.rs b/src/array/sync.rs index bac39b6..140f427 100644 --- a/src/array/sync.rs +++ b/src/array/sync.rs @@ -7,7 +7,7 @@ use crate::data::{for_each_dtype, DataInner, PyData}; use crate::dtype::PyDataType; use crate::error::ZarristaResult; use crate::node::PyNodePath; -use crate::storage::extract_storage; +use crate::storage::PySyncStorage; use ndarray::ArrayD; use pyo3::exceptions::PyNotImplementedError; use pyo3::prelude::*; @@ -44,9 +44,8 @@ impl PyArray { signature = (store, path = PyNodePath::root()), text_signature = "(store, path='/')" )] - fn open(store: &Bound<'_, PyAny>, path: PyNodePath) -> ZarristaResult { - let storage = extract_storage(store)?; - let inner = Array::open(storage, path.as_str())?; + fn open(store: PySyncStorage, path: PyNodePath) -> ZarristaResult { + let inner = Array::open(store.into(), path.as_str())?; Ok(Self::new(inner)) } diff --git a/src/group/sync.rs b/src/group/sync.rs index a8e7ef3..bfaffb0 100644 --- a/src/group/sync.rs +++ b/src/group/sync.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use super::last_segment; use crate::error::ZarristaResult; use crate::node::{open_node, Node, PyNodePath}; -use crate::storage::extract_storage; +use crate::storage::PySyncStorage; use pyo3::prelude::*; use pythonize::pythonize; use pythonize::Result as PythonizeResult; @@ -43,10 +43,10 @@ impl PyGroup { signature = (store, path = PyNodePath::root()), text_signature = "(store, path='/')" )] - fn open(store: &Bound<'_, PyAny>, path: PyNodePath) -> ZarristaResult { - let storage = extract_storage(store)?; - let inner = Group::open(storage.clone(), path.as_str())?; - Ok(Self::new(storage, path.into(), inner)) + fn open(store: PySyncStorage, path: PyNodePath) -> ZarristaResult { + let store: Arc = store.into(); + let inner = Group::open(store.clone(), path.as_str())?; + Ok(Self::new(store, path.into(), inner)) } /// The group's user attributes as a dict. diff --git a/src/storage/mod.rs b/src/storage/mod.rs index ca4aef4..c664e4d 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -6,6 +6,4 @@ mod sync; pub(crate) use python::PyStore; #[allow(unused)] pub use r#async::AsyncStorage; -pub(crate) use sync::extract_storage; -#[allow(unused)] -pub use sync::{PyFilesystemStore, PyMemoryStore}; +pub use sync::{PyFilesystemStore, PyMemoryStore, PySyncStorage}; diff --git a/src/storage/sync.rs b/src/storage/sync.rs index 0a1404d..5d574eb 100644 --- a/src/storage/sync.rs +++ b/src/storage/sync.rs @@ -7,6 +7,30 @@ use zarrs::filesystem::FilesystemStore; use zarrs::storage::store::MemoryStore; use zarrs::storage::ReadableListableStorageTraits; +pub struct PySyncStorage(Arc); + +impl FromPyObject<'_, '_> for PySyncStorage { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, '_, PyAny>) -> Result { + if let Ok(s) = obj.cast::() { + return Ok(Self(s.get().storage.clone())); + } + if let Ok(s) = obj.cast::() { + return Ok(Self(s.get().storage.clone())); + } + + // Any other object is treated as a duck-typed custom store. + Ok(Self(Arc::new(PyStore::new(&obj)))) + } +} + +impl From for Arc { + fn from(s: PySyncStorage) -> Self { + s.0 + } +} + /// A store backed by a local directory. #[pyclass(module = "zarrista", frozen, name = "FilesystemStore")] pub struct PyFilesystemStore { @@ -48,17 +72,3 @@ impl PyMemoryStore { "MemoryStore()".to_string() } } - -/// Pull the inner [`Storage`] out of any zarrista store object. -pub(crate) fn extract_storage( - store: &Bound<'_, PyAny>, -) -> PyResult> { - if let Ok(s) = store.cast::() { - return Ok(s.get().storage.clone()); - } - if let Ok(s) = store.cast::() { - return Ok(s.get().storage.clone()); - } - // Any other object is treated as a duck-typed custom store. - Ok(Arc::new(PyStore::new(store))) -} From 4d14d26aae2e2dbc9e4cd2e553c02bf8bc9a1fde Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Wed, 17 Jun 2026 00:14:52 -0400 Subject: [PATCH 06/10] rename Duck typed storage --- src/storage/mod.rs | 2 +- src/storage/python.rs | 12 ++++++------ src/storage/sync.rs | 4 ++-- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/storage/mod.rs b/src/storage/mod.rs index c664e4d..0f86d16 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -3,7 +3,7 @@ mod key; mod python; mod sync; -pub(crate) use python::PyStore; +pub(crate) use python::PyDuckStore; #[allow(unused)] pub use r#async::AsyncStorage; pub use sync::{PyFilesystemStore, PyMemoryStore, PySyncStorage}; diff --git a/src/storage/python.rs b/src/storage/python.rs index 3ce9b99..dc46863 100644 --- a/src/storage/python.rs +++ b/src/storage/python.rs @@ -1,7 +1,7 @@ //! A custom, duck-typed Python object adapted to the `zarrs` sync storage traits. //! //! The Python object declares capabilities via `@property` predicates and -//! implements a small set of methods. [`PyStore`] reads the capability flags +//! implements a small set of methods. [`PyDuckStore`] reads the capability flags //! once at construction and adapts the object to [`ReadableStorageTraits`] //! (+ [`ListableStorageTraits`], added in a later task). @@ -15,13 +15,13 @@ use zarrs::storage::{ /// A Python object adapted to the `zarrs` sync storage traits. #[derive(Debug)] -pub(crate) struct PyStore { +pub(crate) struct PyDuckStore { obj: Py, supports_get_partial: bool, supports_listing: bool, } -impl PyStore { +impl PyDuckStore { /// Wrap a duck-typed Python store object, reading its capability flags now. pub(crate) fn new(obj: &Bound<'_, PyAny>) -> Self { Self { @@ -64,7 +64,7 @@ fn py_to_storage_error(err: PyErr) -> StorageError { StorageError::Other(err.to_string()) } -impl ReadableStorageTraits for PyStore { +impl ReadableStorageTraits for PyDuckStore { fn get(&self, key: &StoreKey) -> Result { self.py_get(key) } @@ -126,7 +126,7 @@ impl ReadableStorageTraits for PyStore { } } -impl PyStore { +impl PyDuckStore { /// Error returned by every listable method when listing is not declared. fn require_listing(&self) -> Result<(), StorageError> { if self.supports_listing { @@ -159,7 +159,7 @@ impl PyStore { } } -impl ListableStorageTraits for PyStore { +impl ListableStorageTraits for PyDuckStore { fn list(&self) -> Result { self.require_listing()?; Python::attach(|py| self.py_list_keys(py, "list", ())) diff --git a/src/storage/sync.rs b/src/storage/sync.rs index 5d574eb..5284417 100644 --- a/src/storage/sync.rs +++ b/src/storage/sync.rs @@ -1,5 +1,5 @@ use crate::error::ZarristaResult; -use crate::storage::PyStore; +use crate::storage::PyDuckStore; use pyo3::prelude::*; use std::path::PathBuf; use std::sync::Arc; @@ -21,7 +21,7 @@ impl FromPyObject<'_, '_> for PySyncStorage { } // Any other object is treated as a duck-typed custom store. - Ok(Self(Arc::new(PyStore::new(&obj)))) + Ok(Self(Arc::new(PyDuckStore::new(&obj)))) } } From 3e6acfff26376aa8ae79c5b4280b26a00cde1b43 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Wed, 17 Jun 2026 00:17:10 -0400 Subject: [PATCH 07/10] test: custom store listing works when supported, raises when not Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/test_custom_store.py | 44 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 42 insertions(+), 2 deletions(-) diff --git a/tests/test_custom_store.py b/tests/test_custom_store.py index be2354e..b13785d 100644 --- a/tests/test_custom_store.py +++ b/tests/test_custom_store.py @@ -1,10 +1,11 @@ +from __future__ import annotations + import os import numpy as np import pytest import zarr - -from zarrista import Array +from zarrista import Array, Group @pytest.fixture @@ -44,3 +45,42 @@ def test_open_array_and_read_chunk_from_custom_store(zarr_bytes): data = array.retrieve_chunk([0]) assert array.shape == [4] np.testing.assert_array_equal(data.to_numpy(), np.array([0, 1], dtype="int32")) + + +class DictStore(ReadOnlyDictStore): + """Readable + listable dict store.""" + + supports_listing = True + + def list(self) -> list[str]: + return list(self._mapping) + + def list_prefix(self, prefix: str) -> list[str]: + return [k for k in self._mapping if k.startswith(prefix)] + + def list_dir(self, prefix: str) -> dict[str, list[str]]: + keys: list[str] = [] + prefixes: set[str] = set() + for k in self._mapping: + if not k.startswith(prefix): + continue + rest = k[len(prefix):] + if "/" in rest: + prefixes.add(prefix + rest.split("/", 1)[0] + "/") + else: + keys.append(k) + return {"keys": keys, "prefixes": sorted(prefixes)} + + def size_prefix(self, prefix: str) -> int: + return sum(len(v) for k, v in self._mapping.items() if k.startswith(prefix)) + + +def test_listing_works_when_supported(zarr_bytes): + group = Group.open(DictStore(zarr_bytes), "/") + assert group.array_keys() == ["a"] + + +def test_listing_raises_when_unsupported(zarr_bytes): + group = Group.open(ReadOnlyDictStore(zarr_bytes), "/") + with pytest.raises(Exception, match="does not support listing"): + group.array_keys() From 4fe0257a590dc6ff2f5d9073b16a1b74492dcb0b Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Wed, 17 Jun 2026 00:22:44 -0400 Subject: [PATCH 08/10] test: store exceptions surface with their message Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/test_custom_store.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/test_custom_store.py b/tests/test_custom_store.py index b13785d..ed8a963 100644 --- a/tests/test_custom_store.py +++ b/tests/test_custom_store.py @@ -84,3 +84,13 @@ def test_listing_raises_when_unsupported(zarr_bytes): group = Group.open(ReadOnlyDictStore(zarr_bytes), "/") with pytest.raises(Exception, match="does not support listing"): group.array_keys() + + +class BrokenStore(ReadOnlyDictStore): + def get(self, key: str) -> bytes | None: + raise RuntimeError("boom from store") + + +def test_store_exception_surfaces_message(zarr_bytes): + with pytest.raises(Exception, match="boom from store"): + Array.open(BrokenStore(zarr_bytes), "/a") From 44a3dfe23528caff8e7b45de46654660761f3377 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Wed, 17 Jun 2026 00:27:05 -0400 Subject: [PATCH 09/10] feat: delegate partial reads to custom stores that support them When a store declares supports_get_partial, get_partial_many is forwarded to the Python object with (kind, offset, length) byte-range triples instead of fetching and slicing the whole value. Verified via a sharded-array partial-decode test. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/storage/python.rs | 40 +++++++++++++++++++++- tests/test_custom_store.py | 68 +++++++++++++++++++++++++++++++++----- 2 files changed, 99 insertions(+), 9 deletions(-) diff --git a/src/storage/python.rs b/src/storage/python.rs index dc46863..c0e6d08 100644 --- a/src/storage/python.rs +++ b/src/storage/python.rs @@ -64,6 +64,18 @@ fn py_to_storage_error(err: PyErr) -> StorageError { StorageError::Other(err.to_string()) } +/// Encode a `ByteRange` as the Python `(kind, offset, length)` triple. +/// +/// `kind` is `"start"` (read `length` bytes from `offset`, or to the end when +/// `length` is `None`) or `"suffix"` (read the last `length` bytes; `offset` +/// is unused and reported as `0`). +fn byte_range_to_py(byte_range: &ByteRange) -> (&'static str, u64, Option) { + match byte_range { + ByteRange::FromStart(offset, length) => ("start", *offset, *length), + ByteRange::Suffix(length) => ("suffix", 0, Some(*length)), + } +} + impl ReadableStorageTraits for PyDuckStore { fn get(&self, key: &StoreKey) -> Result { self.py_get(key) @@ -74,7 +86,33 @@ impl ReadableStorageTraits for PyDuckStore { key: &StoreKey, byte_ranges: ByteRangeIterator<'a>, ) -> Result, StorageError> { - // MVP: fetch the full value and slice each range (mirrors MemoryStore). + // When the store declares partial support, delegate to its + // `get_partial_many` instead of fetching the whole value. + if self.supports_get_partial { + let encoded: Vec<(&'static str, u64, Option)> = + byte_ranges.map(|br| byte_range_to_py(&br)).collect(); + let bytes: Option> = + Python::attach(|py| -> Result>, StorageError> { + let result = self + .obj + .bind(py) + .call_method1("get_partial_many", (key.as_str(), encoded)) + .map_err(py_to_storage_error)?; + if result.is_none() { + return Ok(None); + } + let raw = result + .extract::>>() + .map_err(py_to_storage_error)?; + Ok(Some(raw.into_iter().map(Bytes::from).collect())) + })?; + return Ok(bytes.map(|v| { + Box::new(v.into_iter().map(Ok)) + as Box>> + })); + } + + // Fallback: fetch the full value and slice each range (mirrors MemoryStore). let Some(data) = self.py_get(key)? else { return Ok(None); }; diff --git a/tests/test_custom_store.py b/tests/test_custom_store.py index ed8a963..17e0518 100644 --- a/tests/test_custom_store.py +++ b/tests/test_custom_store.py @@ -8,6 +8,18 @@ from zarrista import Array, Group +def _dir_to_mapping(root_dir) -> dict[str, bytes]: + """Read every file under `root_dir` into a {relative-key: bytes} mapping.""" + mapping: dict[str, bytes] = {} + for dirpath, _dirs, files in os.walk(root_dir): + for name in files: + full = os.path.join(dirpath, name) + rel = os.path.relpath(full, root_dir).replace(os.sep, "/") + with open(full, "rb") as fh: + mapping[rel] = fh.read() + return mapping + + @pytest.fixture def zarr_bytes(tmp_path): """Write a tiny v3 array with zarr-python, return {key: bytes}.""" @@ -15,15 +27,19 @@ def zarr_bytes(tmp_path): root = zarr.create_group(store=store) arr = root.create_array("a", shape=(4,), chunks=(2,), dtype="int32") arr[:] = np.arange(4, dtype="int32") + return _dir_to_mapping(tmp_path) - mapping: dict[str, bytes] = {} - for dirpath, _dirs, files in os.walk(tmp_path): - for name in files: - full = os.path.join(dirpath, name) - rel = os.path.relpath(full, tmp_path).replace(os.sep, "/") - with open(full, "rb") as fh: - mapping[rel] = fh.read() - return mapping + +@pytest.fixture +def sharded_zarr_bytes(tmp_path): + """Write a sharded v3 array so reads issue byte-range (partial) requests.""" + store = zarr.storage.LocalStore(str(tmp_path)) + root = zarr.create_group(store=store) + arr = root.create_array( + "a", shape=(4,), chunks=(2,), shards=(4,), dtype="int32" + ) + arr[:] = np.arange(4, dtype="int32") + return _dir_to_mapping(tmp_path) class ReadOnlyDictStore: @@ -94,3 +110,39 @@ def get(self, key: str) -> bytes | None: def test_store_exception_surfaces_message(zarr_bytes): with pytest.raises(Exception, match="boom from store"): Array.open(BrokenStore(zarr_bytes), "/a") + + +class PartialStore(DictStore): + """Listable store that also serves partial reads and records the ranges.""" + + supports_get_partial = True + + def __init__(self, mapping): + super().__init__(mapping) + self.partial_calls: list[tuple] = [] + + def get_partial_many(self, key, ranges): + # ranges: list of (kind, offset, length) + self.partial_calls.append((key, tuple(ranges))) + value = self._mapping.get(key) + if value is None: + return None + out = [] + for kind, offset, length in ranges: + if kind == "suffix": + out.append(value[len(value) - length:]) + elif length is None: + out.append(value[offset:]) + else: + out.append(value[offset:offset + length]) + return out + + +def test_partial_reads_are_delegated(sharded_zarr_bytes): + store = PartialStore(sharded_zarr_bytes) + array = Array.open(store, "/a") + # Read a sub-region of the shard so zarrs partial-decodes it (index + + # inner chunk) via byte-range requests rather than fetching the whole shard. + data = array.retrieve_array_subset((slice(0, 2),)) + np.testing.assert_array_equal(data.to_numpy(), np.array([0, 1], dtype="int32")) + assert store.partial_calls, "expected get_partial_many to be delegated" From a677b4bda08bfcc5932ddef5f1e3e6584675217d Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Wed, 17 Jun 2026 14:18:47 -0400 Subject: [PATCH 10/10] docs: custom store Protocols and README section Add runtime-checkable ReadableStore/ListableStore protocols documenting the custom sync store contract, and a README section showing a minimal dict store. Co-Authored-By: Claude Opus 4.8 (1M context) --- README.md | 38 +++++++++++++++++++++++++ python/zarrista/__init__.py | 3 ++ python/zarrista/_protocols.py | 52 +++++++++++++++++++++++++++++++++++ tests/test_custom_store.py | 7 +++++ 4 files changed, 100 insertions(+) create mode 100644 python/zarrista/_protocols.py diff --git a/README.md b/README.md index cc33690..b52aafe 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,44 @@ A low-level Zarr API for Python, inspired by [zarrita.js], powered from Rust by This has been _minimally_ vibe-coded (Claude still writes bad Rust code in my opinion). +## Custom stores + +Besides the built-in `FilesystemStore` and `MemoryStore`, you can pass any +duck-typed Python object as a (synchronous) store. The minimal contract is a +single `get` method plus two capability properties: + +```python +from zarrista import Array + + +class DictStore: + supports_get_partial = False # opt into byte-range reads + supports_listing = False # opt into listing keys/prefixes + + def __init__(self, mapping: dict[str, bytes]): + self._mapping = mapping + + def get(self, key: str) -> bytes | None: + return self._mapping.get(key) + + +array = Array.open(DictStore(my_bytes), "/path") +``` + +Declare `supports_listing = True` and implement `list`, `list_prefix`, +`list_dir`, and `size_prefix` to support operations like `Group.array_keys()`; +calling a listing operation on a store that does not support it raises an error. +Declare `supports_get_partial = True` and implement `get_partial_many` to serve +efficient byte-range reads (otherwise partial reads fall back to fetching the +whole value and slicing). The `zarrista.ReadableStore` and +`zarrista.ListableStore` protocols document the full surface. + +> Note: if your store defines a method named `list`, add +> `from __future__ import annotations` to the module so later `list[...]` type +> annotations are not shadowed by the method. + +This is sync-only; for async use pass an `obstore.ObjectStore`. + ## Development Requires a Rust toolchain and Python 3.11+. We use diff --git a/python/zarrista/__init__.py b/python/zarrista/__init__.py index 5c56fdc..d39fb15 100644 --- a/python/zarrista/__init__.py +++ b/python/zarrista/__init__.py @@ -1,3 +1,4 @@ +from ._protocols import ListableStore, ReadableStore from ._zarrista import ( Array, AsyncArray, @@ -22,6 +23,8 @@ "DataType", "FilesystemStore", "Group", + "ListableStore", "MemoryStore", + "ReadableStore", "__version__", ] diff --git a/python/zarrista/_protocols.py b/python/zarrista/_protocols.py new file mode 100644 index 0000000..328efa5 --- /dev/null +++ b/python/zarrista/_protocols.py @@ -0,0 +1,52 @@ +"""Typing protocols for custom, duck-typed sync stores. + +A custom store is any Python object satisfying :class:`ReadableStore` (and +optionally :class:`ListableStore`). Pass one anywhere a built-in store is +accepted, e.g. ``Array.open(my_store, "/path")``. + +Capabilities are declared with ``@property`` predicates and read once when the +store is wrapped. ``get`` is the only required method. + +Two methods are *optional* and consulted only when ``supports_get_partial`` is +true (they are intentionally not part of the runtime-checkable protocol, so a +minimal ``get``-only store still satisfies :class:`ReadableStore`): + +- ``get_partial_many(key, ranges) -> list[bytes] | None`` where each range is a + ``(kind, offset, length)`` triple. ``kind`` is ``"start"`` (read ``length`` + bytes from ``offset``, or to the end when ``length`` is ``None``) or + ``"suffix"`` (read the last ``length`` bytes). When absent, partial reads + fall back to fetching the whole value and slicing. +- ``size_key(key) -> int | None``. When absent, the size falls back to + ``len(get(key))``. +""" + +from __future__ import annotations + +import builtins +from typing import Protocol, runtime_checkable + + +@runtime_checkable +class ReadableStore(Protocol): + """A duck-typed, readable sync store. ``get`` is the only required method.""" + + @property + def supports_get_partial(self) -> bool: ... + + @property + def supports_listing(self) -> bool: ... + + def get(self, key: str) -> bytes | None: ... + + +@runtime_checkable +class ListableStore(ReadableStore, Protocol): + """A readable store that also supports listing keys and prefixes.""" + + def list(self) -> builtins.list[str]: ... + + def list_prefix(self, prefix: str) -> builtins.list[str]: ... + + def list_dir(self, prefix: str) -> dict[str, builtins.list[str]]: ... + + def size_prefix(self, prefix: str) -> int: ... diff --git a/tests/test_custom_store.py b/tests/test_custom_store.py index 17e0518..f7aaa5a 100644 --- a/tests/test_custom_store.py +++ b/tests/test_custom_store.py @@ -146,3 +146,10 @@ def test_partial_reads_are_delegated(sharded_zarr_bytes): data = array.retrieve_array_subset((slice(0, 2),)) np.testing.assert_array_equal(data.to_numpy(), np.array([0, 1], dtype="int32")) assert store.partial_calls, "expected get_partial_many to be delegated" + + +def test_store_protocols_are_importable(): + from zarrista import ListableStore, ReadableStore # noqa: F401 + + assert isinstance(ReadOnlyDictStore({}), ReadableStore) + assert isinstance(DictStore({}), ListableStore)