From d836dc1633fc4fb864e45e091233d8852dbd15e2 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Mon, 20 Jul 2026 15:24:28 +0200 Subject: [PATCH 1/7] Feat add sync store attribute --- src/array/builder.rs | 4 ++-- src/array/sync.rs | 39 +++++++++++++++++----------------- src/group/sync.rs | 35 ++++++++++++------------------- src/node/sync.rs | 16 +++++++------- src/storage/sync.rs | 50 +++++++++++++++++++++++++++++--------------- 5 files changed, 75 insertions(+), 69 deletions(-) diff --git a/src/array/builder.rs b/src/array/builder.rs index 660f255..9f36e45 100644 --- a/src/array/builder.rs +++ b/src/array/builder.rs @@ -81,9 +81,9 @@ impl PyArrayBuilder { } fn create(&self, store: PySyncStorage, path: &str) -> ZarristaResult { - let array = self.0.build_arc(store.into_inner(), path)?; + let array = self.0.build_arc(store.inner(), path)?; array.store_metadata()?; - Ok(array.into()) + Ok(PyArray::new(array, store)) } #[cfg(feature = "async")] diff --git a/src/array/sync.rs b/src/array/sync.rs index 58e35a4..53ac913 100644 --- a/src/array/sync.rs +++ b/src/array/sync.rs @@ -21,13 +21,17 @@ use zarrs::storage::ReadableWritableListableStorageTraits; #[pyclass(module = "zarrista", frozen, name = "Array")] pub struct PyArray { pub(crate) inner: Arc>, + pub(crate) store: PySyncStorage, } crate::wasm_send_sync!(PyArray); impl PyArray { - pub(crate) fn new(inner: Arc>) -> Self { - Self { inner } + pub(crate) fn new( + inner: Arc>, + store: PySyncStorage, + ) -> Self { + Self { inner, store } } pub fn inner(&self) -> &Arc> { @@ -66,9 +70,8 @@ impl PyArray { store: PySyncStorage, path: PyNodePath, ) -> ZarristaResult { - let inner = - Array::new_with_metadata(store.into_inner(), path.as_str(), metadata.into_inner())?; - Ok(Self::new(Arc::new(inner))) + let inner = Array::new_with_metadata(store.inner(), path.as_str(), metadata.into_inner())?; + Ok(Self::new(Arc::new(inner), store)) } /// Open the array stored at `path` in `store`. @@ -78,8 +81,8 @@ impl PyArray { text_signature = "(store, path='/')" )] fn open(store: PySyncStorage, path: PyNodePath) -> ZarristaResult { - let inner = Array::open(store.into(), path.as_str())?; - Ok(Self::new(Arc::new(inner))) + let inner = Array::open(store.inner(), path.as_str())?; + Ok(Self::new(Arc::new(inner), store)) } #[pyo3(signature = (chunk_indices, **codec_options))] @@ -107,7 +110,10 @@ impl PyArray { fn read_only(&self) -> Self { let read_list_storage = self.inner.storage().readable_listable(); let storage = Arc::new(ReadOnlyStorageAdapter::new(read_list_storage)); - Self::new(Arc::new(self.inner.with_storage(storage))) + Self::new( + Arc::new(self.inner.with_storage(storage)), + self.store.clone(), + ) } /// Read a region of the array, using numpy-style basic indexing. @@ -172,6 +178,11 @@ impl PyArray { .retrieve_subchunk_opt(&subchunk_cache, &subchunk_indices, &codec_options)?) } + #[getter] + fn store(&self) -> PySyncStorage { + self.store.clone() + } + #[pyo3(signature = (chunk_indices, decoded_chunk, **codec_options))] fn store_chunk( &self, @@ -204,15 +215,3 @@ impl PyArray { Ok(()) } } - -impl From> for PyArray { - fn from(inner: Array) -> Self { - Self::new(Arc::new(inner)) - } -} - -impl From>> for PyArray { - fn from(inner: Arc>) -> Self { - Self::new(inner) - } -} diff --git a/src/group/sync.rs b/src/group/sync.rs index 8613bf6..0b65da3 100644 --- a/src/group/sync.rs +++ b/src/group/sync.rs @@ -20,13 +20,17 @@ use zarrs::storage::ReadableWritableListableStorageTraits; #[pyclass(module = "zarrista", frozen, name = "Group")] pub struct PyGroup { pub(crate) inner: Arc>, + store: PySyncStorage, } crate::wasm_send_sync!(PyGroup); impl PyGroup { - pub(crate) fn new(inner: Arc>) -> Self { - Self { inner } + pub(crate) fn new( + inner: Arc>, + store: PySyncStorage, + ) -> Self { + Self { inner, store } } fn storage(&self) -> Arc { @@ -46,9 +50,8 @@ impl PyGroup { text_signature = "(store, path='/')" )] fn open(store: PySyncStorage, path: PyNodePath) -> ZarristaResult { - let store = store.into_inner(); - let inner = Group::open(store, path.as_str())?; - Ok(Self::new(Arc::new(inner))) + let inner = Group::open(store.inner(), path.as_str())?; + Ok(Self::new(Arc::new(inner), store)) } /// Names of the direct child arrays. @@ -74,7 +77,7 @@ impl PyGroup { .into_iter() .find(|child| child.name().as_str() == name.as_str()) .ok_or(PyKeyError::new_err(format!("child {name} not found")))?; - Ok(PyNode::new(selected_child, self.storage())) + Ok(PyNode::new(selected_child, self.store.clone())) } /// Every node under the group, recursively, as `Array`/`Group` objects. @@ -88,12 +91,12 @@ impl PyGroup { NodeMetadata::Array(array_metadata) => { let array = Array::new_with_metadata(storage, path.as_str(), array_metadata)?; - Ok(PyArray::new(Arc::new(array)).into()) + Ok(PyArray::new(Arc::new(array), self.store.clone()).into()) } NodeMetadata::Group(group_metadata) => { let group = Group::new_with_metadata(storage, path.as_str(), group_metadata)?; - Ok(PyGroup::new(Arc::new(group)).into()) + Ok(PyGroup::new(Arc::new(group), self.store.clone()).into()) } } }) @@ -106,7 +109,7 @@ impl PyGroup { .inner .child_arrays()? .into_iter() - .map(|array| array.into()) + .map(|array| PyArray::new(Arc::new(array), self.store.clone())) .collect()) } @@ -116,7 +119,7 @@ impl PyGroup { .inner .child_groups()? .into_iter() - .map(Self::from) + .map(|group| PyGroup::new(Arc::new(group), self.store.clone())) .collect()) } @@ -166,15 +169,3 @@ impl PyGroup { format!("Group(path={:?})", self.inner.path().as_str()) } } - -impl From> for PyGroup { - fn from(inner: Group) -> Self { - Self::new(Arc::new(inner)) - } -} - -impl From>> for PyGroup { - fn from(inner: Arc>) -> Self { - Self::new(inner) - } -} diff --git a/src/node/sync.rs b/src/node/sync.rs index 16cad39..7110451 100644 --- a/src/node/sync.rs +++ b/src/node/sync.rs @@ -5,21 +5,21 @@ use std::sync::Arc; use crate::array::PyArray; use crate::error::ZarristaError; use crate::group::PyGroup; +use crate::storage::PySyncStorage; use pyo3::IntoPyObjectExt; use pyo3::prelude::*; use zarrs::array::Array; use zarrs::group::Group; use zarrs::node::{Node, NodeMetadata}; -use zarrs::storage::ReadableWritableListableStorageTraits; /// An opened node: either an array or a group. pub(crate) struct PyNode { node: Node, - storage: Arc, + storage: PySyncStorage, } impl PyNode { - pub fn new(node: Node, storage: Arc) -> Self { + pub fn new(node: Node, storage: PySyncStorage) -> Self { Self { node, storage } } } @@ -31,17 +31,17 @@ impl<'py> IntoPyObject<'py> for PyNode { type Output = Bound<'py, Self::Target>; fn into_pyobject(self, py: Python<'py>) -> Result { - let storage = self.storage; + let store = self.storage; let path = self.node.path().clone(); let node_metadata = NodeMetadata::from(self.node); match node_metadata { NodeMetadata::Array(array_metadata) => { - let array = Array::new_with_metadata(storage, path.as_str(), array_metadata)?; - Ok(PyArray::new(Arc::new(array)).into_bound_py_any(py)?) + let array = Array::new_with_metadata(store.inner(), path.as_str(), array_metadata)?; + Ok(PyArray::new(Arc::new(array), store).into_bound_py_any(py)?) } NodeMetadata::Group(group_metadata) => { - let group = Group::new_with_metadata(storage, path.as_str(), group_metadata)?; - Ok(PyGroup::new(Arc::new(group)).into_bound_py_any(py)?) + let group = Group::new_with_metadata(store.inner(), path.as_str(), group_metadata)?; + Ok(PyGroup::new(Arc::new(group), store).into_bound_py_any(py)?) } } } diff --git a/src/storage/sync.rs b/src/storage/sync.rs index 9356590..2f0eb00 100644 --- a/src/storage/sync.rs +++ b/src/storage/sync.rs @@ -11,12 +11,20 @@ use zarrs::storage::{ ReadableListableStorageTraits, ReadableStorageTraits, ReadableWritableListableStorageTraits, StorageError, StoreKey, StoreKeys, StoreKeysPrefixes, StorePrefix, WritableStorageTraits, }; + /// A zarrista sync store object adapted to the maximal `zarrs` storage trait. -pub struct PySyncStorage(Arc); +#[derive(Clone, IntoPyObject)] +pub enum PySyncStorage { + Filesystem(PyFilesystemStore), + MemoryStore(PyMemoryStore), +} impl PySyncStorage { - pub fn into_inner(self) -> Arc { - self.0 + pub fn inner(&self) -> Arc { + match self { + Self::Filesystem(store) => store.storage.clone(), + Self::MemoryStore(store) => store.0.clone(), + } } } @@ -25,10 +33,10 @@ impl FromPyObject<'_, '_> for PySyncStorage { fn extract(obj: Borrowed<'_, '_, PyAny>) -> Result { if let Ok(s) = obj.cast::() { - return Ok(Self(s.get().storage.clone())); + return Ok(Self::Filesystem(s.get().clone())); } if let Ok(s) = obj.cast::() { - return Ok(Self(s.get().storage.clone())); + return Ok(Self::MemoryStore(s.get().clone())); } Err(PyTypeError::new_err( "expected a FilesystemStore or MemoryStore", @@ -38,14 +46,24 @@ impl FromPyObject<'_, '_> for PySyncStorage { impl From for Arc { fn from(s: PySyncStorage) -> Self { - s.0 + match s { + PySyncStorage::Filesystem(store) => store.storage, + PySyncStorage::MemoryStore(store) => store.0, + } } } /// A store backed by a local directory. -#[pyclass(module = "zarrista", frozen, name = "FilesystemStore")] +#[pyclass( + module = "zarrista", + frozen, + name = "FilesystemStore", + skip_from_py_object +)] +#[derive(Clone)] pub struct PyFilesystemStore { - pub(crate) storage: Arc, + pub(crate) storage: Arc, + path: PathBuf, } crate::wasm_send_sync!(PyFilesystemStore); @@ -55,22 +73,22 @@ impl PyFilesystemStore { /// Open a filesystem store rooted at `path`. #[new] fn new(path: PathBuf) -> ZarristaResult { - let store = FilesystemStore::new(path)?; + let store = FilesystemStore::new(&path)?; Ok(Self { storage: Arc::new(store), + path, }) } fn __repr__(&self) -> String { - "FilesystemStore(...)".to_string() + format!("FilesystemStore({})", self.path.display()) } } /// An in-memory store, primarily useful for testing. -#[pyclass(module = "zarrista", frozen, name = "MemoryStore")] -pub struct PyMemoryStore { - pub(crate) storage: Arc, -} +#[pyclass(module = "zarrista", frozen, name = "MemoryStore", skip_from_py_object)] +#[derive(Clone)] +pub struct PyMemoryStore(Arc); crate::wasm_send_sync!(PyMemoryStore); @@ -78,9 +96,7 @@ crate::wasm_send_sync!(PyMemoryStore); impl PyMemoryStore { #[new] fn new() -> Self { - Self { - storage: Arc::new(MemoryStore::new()), - } + Self(Arc::new(MemoryStore::new())) } fn __repr__(&self) -> String { From 8f9d18901b78df9821dad9016d048cb821fea809 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Mon, 20 Jul 2026 16:01:48 +0200 Subject: [PATCH 2/7] add async storage export to python --- src/array/async.rs | 40 +++++++++--------- src/array/builder.rs | 4 +- src/group/async.rs | 71 ++++++++++++++++++-------------- src/node/async.rs | 19 ++++----- src/storage/async.rs | 98 ++++++++++++++++++++++++++++++++++++-------- 5 files changed, 150 insertions(+), 82 deletions(-) diff --git a/src/array/async.rs b/src/array/async.rs index e8860f1..9144cef 100644 --- a/src/array/async.rs +++ b/src/array/async.rs @@ -23,11 +23,15 @@ use zarrs::storage::AsyncReadableWritableListableStorageTraits; #[pyclass(module = "zarrista", frozen, name = "AsyncArray", from_py_object)] pub struct PyAsyncArray { pub(crate) inner: Arc>, + store: PyAsyncStorage, } impl PyAsyncArray { - pub(crate) fn new(inner: Arc>) -> Self { - Self { inner } + pub(crate) fn new( + inner: Arc>, + store: PyAsyncStorage, + ) -> Self { + Self { inner, store } } pub fn inner(&self) -> &Arc> { @@ -70,9 +74,8 @@ impl PyAsyncArray { store: PyAsyncStorage, path: PyNodePath, ) -> ZarristaResult { - let inner = - Array::new_with_metadata(store.into_inner(), path.as_str(), metadata.into_inner())?; - Ok(Self::new(Arc::new(inner))) + let inner = Array::new_with_metadata(store.inner(), path.as_str(), metadata.into_inner())?; + Ok(Self::new(Arc::new(inner), store)) } /// Open the array stored at `path` in `store`. @@ -86,12 +89,11 @@ impl PyAsyncArray { store: PyAsyncStorage, path: PyNodePath, ) -> PyResult> { - let storage = store.into(); future_into_py(py, async move { - let inner = Array::async_open(storage, path.as_str()) + let inner = Array::async_open(store.inner(), path.as_str()) .await .map_err(ZarristaError::from)?; - Ok(Self::new(Arc::new(inner))) + Ok(Self::new(Arc::new(inner), store)) }) } @@ -135,7 +137,10 @@ impl PyAsyncArray { fn read_only(&self) -> Self { let read_list_storage = self.inner.storage().readable_listable(); let storage = Arc::new(AsyncReadOnlyStorageAdapter::new(read_list_storage)); - Self::new(Arc::new(self.inner.with_storage(storage))) + Self::new( + Arc::new(self.inner.with_storage(storage)), + self.store.clone(), + ) } fn erase_metadata<'py>(&self, py: Python<'py>) -> PyResult> { @@ -249,6 +254,11 @@ impl PyAsyncArray { }) } + #[getter] + fn store(&self) -> &PyAsyncStorage { + &self.store + } + #[pyo3(signature = (chunk_indices, decoded_chunk, **codec_options))] fn store_chunk<'py>( &self, @@ -295,15 +305,3 @@ impl PyAsyncArray { }) } } - -impl From> for PyAsyncArray { - fn from(inner: Array) -> Self { - Self::new(Arc::new(inner)) - } -} - -impl From>> for PyAsyncArray { - fn from(inner: Arc>) -> Self { - Self::new(inner) - } -} diff --git a/src/array/builder.rs b/src/array/builder.rs index 9f36e45..4dba5ab 100644 --- a/src/array/builder.rs +++ b/src/array/builder.rs @@ -97,7 +97,7 @@ impl PyArrayBuilder { let array = self .0 - .build_arc(store.into_inner(), path) + .build_arc(store.inner(), path) .map_err(ZarristaError::from)?; pyo3_async_runtimes::tokio::future_into_py(py, async move { @@ -105,7 +105,7 @@ impl PyArrayBuilder { .async_store_metadata() .await .map_err(ZarristaError::from)?; - Ok(crate::array::PyAsyncArray::from(array)) + Ok(crate::array::PyAsyncArray::new(array, store)) }) } diff --git a/src/group/async.rs b/src/group/async.rs index 59fddcc..3c5f5cd 100644 --- a/src/group/async.rs +++ b/src/group/async.rs @@ -19,15 +19,15 @@ use zarrs::storage::AsyncReadableWritableListableStorageTraits; #[pyclass(module = "zarrista", frozen, name = "AsyncGroup", from_py_object)] pub struct PyAsyncGroup { pub(crate) inner: Arc>, + store: PyAsyncStorage, } impl PyAsyncGroup { - pub(crate) fn new(inner: Arc>) -> Self { - Self { inner } - } - - fn storage(&self) -> Arc { - self.inner.storage() + pub(crate) fn new( + inner: Arc>, + store: PyAsyncStorage, + ) -> Self { + Self { inner, store } } } @@ -47,12 +47,11 @@ impl PyAsyncGroup { store: PyAsyncStorage, path: PyNodePath, ) -> PyResult> { - let storage = store.into_inner(); future_into_py(py, async move { - let inner = Group::async_open(storage.clone(), path.as_str()) + let inner = Group::async_open(store.inner(), path.as_str()) .await .map_err(ZarristaError::from)?; - Ok(Self::new(Arc::new(inner))) + Ok(Self::new(Arc::new(inner), store)) }) } @@ -89,7 +88,7 @@ impl PyAsyncGroup { /// Open a direct child array or group by name. fn open_child_async<'py>(&self, py: Python<'py>, name: String) -> PyResult> { let inner = self.inner.clone(); - let storage = self.storage(); + let storage = self.store.clone(); future_into_py(py, async move { let children = inner .async_children(false) @@ -105,7 +104,7 @@ impl PyAsyncGroup { /// Every node under the group, recursively, as `AsyncArray`/`AsyncGroup` objects. fn traverse<'py>(&self, py: Python<'py>) -> PyResult> { - let storage = self.storage(); + let storage = self.store.clone(); let inner = self.inner.clone(); future_into_py(py, async move { let nodes = inner.async_traverse().await.map_err(ZarristaError::from)?; @@ -115,14 +114,20 @@ impl PyAsyncGroup { let storage = storage.clone(); match metadata { NodeMetadata::Array(array_metadata) => { - let array = - Array::new_with_metadata(storage, path.as_str(), array_metadata)?; - Ok(PyAsyncArray::new(Arc::new(array)).into()) + let array = Array::new_with_metadata( + storage.inner(), + path.as_str(), + array_metadata, + )?; + Ok(PyAsyncArray::new(Arc::new(array), storage.clone()).into()) } NodeMetadata::Group(group_metadata) => { - let group = - Group::new_with_metadata(storage, path.as_str(), group_metadata)?; - Ok(PyAsyncGroup::new(Arc::new(group)).into()) + let group = Group::new_with_metadata( + storage.inner(), + path.as_str(), + group_metadata, + )?; + Ok(PyAsyncGroup::new(Arc::new(group), storage.clone()).into()) } } }) @@ -134,6 +139,7 @@ impl PyAsyncGroup { /// The direct child arrays of the group. fn child_arrays<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); + let store = self.store.clone(); future_into_py(py, async move { let arrays = inner .async_child_arrays() @@ -141,7 +147,7 @@ impl PyAsyncGroup { .map_err(ZarristaError::from)?; Ok(arrays .into_iter() - .map(|array| PyAsyncArray::new(Arc::new(array))) + .map(|array| PyAsyncArray::new(Arc::new(array), store.clone())) .collect::>()) }) } @@ -149,12 +155,16 @@ impl PyAsyncGroup { /// The direct child groups of the group. fn child_groups<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); + let store = self.store.clone(); future_into_py(py, async move { let groups = inner .async_child_groups() .await .map_err(ZarristaError::from)?; - Ok(groups.into_iter().map(Self::from).collect::>()) + Ok(groups + .into_iter() + .map(|group| PyAsyncGroup::new(Arc::new(group), store.clone())) + .collect::>()) }) } @@ -194,24 +204,29 @@ impl PyAsyncGroup { }) } - /// Write the group metadata to the store. - fn store_metadata<'py>(&self, py: Python<'py>) -> PyResult> { + /// Erase the group metadata from the store. Succeeds if it does not exist. + fn erase_metadata<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); future_into_py(py, async move { inner - .async_store_metadata() + .async_erase_metadata() .await .map_err(ZarristaError::from)?; Ok(()) }) } - /// Erase the group metadata from the store. Succeeds if it does not exist. - fn erase_metadata<'py>(&self, py: Python<'py>) -> PyResult> { + #[getter] + fn store(&self) -> &PyAsyncStorage { + &self.store + } + + /// Write the group metadata to the store. + fn store_metadata<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); future_into_py(py, async move { inner - .async_erase_metadata() + .async_store_metadata() .await .map_err(ZarristaError::from)?; Ok(()) @@ -222,9 +237,3 @@ impl PyAsyncGroup { format!("AsyncGroup(path={:?})", self.inner.path().as_str()) } } - -impl From> for PyAsyncGroup { - fn from(inner: Group) -> Self { - Self::new(Arc::new(inner)) - } -} diff --git a/src/node/async.rs b/src/node/async.rs index d32dd71..b26e830 100644 --- a/src/node/async.rs +++ b/src/node/async.rs @@ -5,24 +5,21 @@ use std::sync::Arc; use crate::array::PyAsyncArray; use crate::error::{ZarristaError, ZarristaResult}; use crate::group::PyAsyncGroup; +use crate::storage::PyAsyncStorage; use pyo3::IntoPyObjectExt; use pyo3::prelude::*; use zarrs::array::Array; use zarrs::group::Group; use zarrs::node::{Node, NodeMetadata}; -use zarrs::storage::AsyncReadableWritableListableStorageTraits; /// An opened node from an async store: either an array or a group. pub(crate) struct PyAsyncNode { node: Node, - storage: Arc, + storage: PyAsyncStorage, } impl PyAsyncNode { - pub fn new( - node: Node, - storage: Arc, - ) -> ZarristaResult { + pub fn new(node: Node, storage: PyAsyncStorage) -> ZarristaResult { Ok(Self { node, storage }) } } @@ -39,12 +36,14 @@ impl<'py> IntoPyObject<'py> for PyAsyncNode { let node_metadata = NodeMetadata::from(self.node); match node_metadata { NodeMetadata::Array(array_metadata) => { - let array = Array::new_with_metadata(storage, path.as_str(), array_metadata)?; - Ok(PyAsyncArray::new(Arc::new(array)).into_bound_py_any(py)?) + let array = + Array::new_with_metadata(storage.inner(), path.as_str(), array_metadata)?; + Ok(PyAsyncArray::new(Arc::new(array), storage).into_bound_py_any(py)?) } NodeMetadata::Group(group_metadata) => { - let group = Group::new_with_metadata(storage, path.as_str(), group_metadata)?; - Ok(PyAsyncGroup::new(Arc::new(group)).into_bound_py_any(py)?) + let group = + Group::new_with_metadata(storage.inner(), path.as_str(), group_metadata)?; + Ok(PyAsyncGroup::new(Arc::new(group), storage).into_bound_py_any(py)?) } } } diff --git a/src/storage/async.rs b/src/storage/async.rs index 2661fc2..51e5f01 100644 --- a/src/storage/async.rs +++ b/src/storage/async.rs @@ -1,3 +1,4 @@ +use std::convert::Infallible; use std::sync::Arc; use async_trait::async_trait; @@ -17,11 +18,18 @@ use zarrs::storage::{ use zarrs_icechunk::AsyncIcechunkStore; use zarrs_object_store::AsyncObjectStore; -pub struct PyAsyncStorage(Arc); +#[derive(Clone, IntoPyObject)] +pub enum PyAsyncStorage { + ObjectStore(PyAsyncObjectStore), + Icechunk(PyAsyncIcechunkStore), +} impl PyAsyncStorage { - pub fn into_inner(self) -> Arc { - self.0 + pub fn inner(&self) -> Arc { + match self { + Self::ObjectStore(store) => store.inner.clone(), + Self::Icechunk(store) => store.inner.clone(), + } } } @@ -33,11 +41,11 @@ impl FromPyObject<'_, '_> for PyAsyncStorage { // falling through to the generic error below: the caller meant icechunk. if is_icechunk_session(obj)? { let store = obj.extract::()?; - return Ok(Self(Arc::new(store.into_inner()))); + return Ok(Self::Icechunk(store)); } if let Ok(store) = obj.extract::() { - return Ok(Self(Arc::new(store.into_inner()))); + return Ok(Self::ObjectStore(store)); } Err(PyTypeError::new_err( @@ -46,18 +54,29 @@ impl FromPyObject<'_, '_> for PyAsyncStorage { } } +impl<'py> IntoPyObject<'py> for &PyAsyncStorage { + type Target = PyAny; + type Error = Infallible; + type Output = Bound<'py, Self::Target>; + + fn into_pyobject(self, py: Python<'py>) -> Result { + match self { + PyAsyncStorage::ObjectStore(store) => store.into_pyobject(py), + PyAsyncStorage::Icechunk(store) => store.into_pyobject(py), + } + } +} + impl From for Arc { fn from(s: PyAsyncStorage) -> Self { - s.0 + s.inner() } } -pub struct PyAsyncObjectStore(AsyncObjectStore>); - -impl PyAsyncObjectStore { - fn into_inner(self) -> AsyncObjectStore> { - self.0 - } +#[derive(Clone)] +pub struct PyAsyncObjectStore { + inner: Arc>>, + pyobj: Arc>, } impl FromPyObject<'_, '_> for PyAsyncObjectStore { @@ -65,18 +84,39 @@ impl FromPyObject<'_, '_> for PyAsyncObjectStore { fn extract(obj: Borrowed<'_, '_, PyAny>) -> Result { let store = obj.extract::()?; - Ok(Self(AsyncObjectStore::new(store.into_dyn()))) + Ok(Self { + inner: Arc::new(AsyncObjectStore::new(store.into_dyn())), + pyobj: Arc::new(obj.into()), + }) + } +} + +impl<'py> IntoPyObject<'py> for PyAsyncObjectStore { + type Target = PyAny; + type Error = Infallible; + type Output = Bound<'py, Self::Target>; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(self.pyobj.bind(py).clone()) } } -pub struct PyAsyncIcechunkStore(AsyncIcechunkStore); +impl<'py> IntoPyObject<'py> for &PyAsyncObjectStore { + type Target = PyAny; + type Error = Infallible; + type Output = Bound<'py, Self::Target>; -impl PyAsyncIcechunkStore { - fn into_inner(self) -> AsyncIcechunkStore { - self.0 + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(self.pyobj.bind(py).clone()) } } +#[derive(Clone)] +pub struct PyAsyncIcechunkStore { + inner: Arc, + pyobj: Arc>, +} + /// Whether `obj` is an `icechunk.session.Session` instance. fn is_icechunk_session(obj: Borrowed<'_, '_, PyAny>) -> PyResult { let class = obj.getattr("__class__")?; @@ -127,10 +167,32 @@ impl FromPyObject<'_, '_> for PyAsyncIcechunkStore { against (2.x)." )) })?; - Ok(Self(AsyncIcechunkStore::new(session))) + Ok(Self { + inner: Arc::new(AsyncIcechunkStore::new(session)), + pyobj: Arc::new(obj.into()), + }) } } +impl<'py> IntoPyObject<'py> for PyAsyncIcechunkStore { + type Target = PyAny; + type Error = Infallible; + type Output = Bound<'py, Self::Target>; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(self.pyobj.bind(py).clone()) + } +} + +impl<'py> IntoPyObject<'py> for &PyAsyncIcechunkStore { + type Target = PyAny; + type Error = Infallible; + type Output = Bound<'py, Self::Target>; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(self.pyobj.bind(py).clone()) + } +} /// An async storage adapter that reads and lists transparently but rejects all writes at runtime. pub struct AsyncReadOnlyStorageAdapter(Arc); From 031c7efc779d0cfef871098ad1d3155e0e1c724c Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Mon, 20 Jul 2026 16:02:14 +0200 Subject: [PATCH 3/7] return store from group --- src/group/sync.rs | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/src/group/sync.rs b/src/group/sync.rs index 0b65da3..188301f 100644 --- a/src/group/sync.rs +++ b/src/group/sync.rs @@ -153,18 +153,23 @@ impl PyGroup { .collect()) } - /// Write the group metadata to the store. - fn store_metadata(&self) -> ZarristaResult<()> { - self.inner.store_metadata()?; - Ok(()) - } - /// Erase the group metadata from the store. Succeeds if it does not exist. fn erase_metadata(&self) -> ZarristaResult<()> { self.inner.erase_metadata()?; Ok(()) } + #[getter] + fn store(&self) -> PySyncStorage { + self.store.clone() + } + + /// Write the group metadata to the store. + fn store_metadata(&self) -> ZarristaResult<()> { + self.inner.store_metadata()?; + Ok(()) + } + fn __repr__(&self) -> String { format!("Group(path={:?})", self.inner.path().as_str()) } From 37193ade848be6fbd1e9461e166c41267382a494 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Mon, 20 Jul 2026 16:10:03 +0200 Subject: [PATCH 4/7] Add python type hints --- python/zarrista/_array.pyi | 6 ++++++ python/zarrista/_group.pyi | 10 ++++++++-- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/python/zarrista/_array.pyi b/python/zarrista/_array.pyi index a636c75..66a16d9 100644 --- a/python/zarrista/_array.pyi +++ b/python/zarrista/_array.pyi @@ -149,6 +149,9 @@ class Array: The bytes are returned verbatim, without running the codec pipeline. Returns `None` if the subchunk is absent from the store. """ + @property + def store(self) -> FilesystemStore | MemoryStore: + """Retrieve the store backing this array.""" def store_chunk( self, chunk_indices: list[int], @@ -369,6 +372,9 @@ class AsyncArray: The bytes are returned verbatim, without running the codec pipeline. Returns `None` if the subchunk is absent from the store. """ + @property + def store(self) -> AsyncStore: + """Retrieve the store backing this array.""" async def store_chunk( self, chunk_indices: list[int], diff --git a/python/zarrista/_group.pyi b/python/zarrista/_group.pyi index 38de785..b6b7a4f 100644 --- a/python/zarrista/_group.pyi +++ b/python/zarrista/_group.pyi @@ -37,12 +37,15 @@ class Group: """Return the full paths of the group's direct child arrays.""" def child_group_paths(self) -> list[str]: """Return the full paths of the group's direct child groups.""" - def store_metadata(self) -> None: - """Write the group metadata to the store.""" def erase_metadata(self) -> None: """Erase the group metadata from the store. Succeeds if it does not exist.""" def child(self, name: str) -> Array | Group: """Open a direct child array or group by name.""" + def store_metadata(self) -> None: + """Write the group metadata to the store.""" + @property + def store(self) -> FilesystemStore | MemoryStore: + """Retrieve the store backing this group.""" def __getitem__(self, name: str) -> Array | Group: """Open a direct child array or group by name.""" @@ -89,3 +92,6 @@ class AsyncGroup: """Erase the group metadata from the store. Succeeds if it does not exist.""" async def open_child_async(self, name: str) -> AsyncArray | AsyncGroup: """Open a direct child array or group by name.""" + @property + def store(self) -> AsyncStore: + """Retrieve the store backing this group.""" From 2295012279076b5aa12b0062dc883d620080b1cc Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Mon, 20 Jul 2026 16:10:22 +0200 Subject: [PATCH 5/7] remove dev docs from front readme --- README.md | 21 --------------------- 1 file changed, 21 deletions(-) diff --git a/README.md b/README.md index 22ecb1d..65f1ca3 100644 --- a/README.md +++ b/README.md @@ -77,24 +77,3 @@ You can also read individual chunks by their grid index: ```py data = array.retrieve_chunk([0, 0]) ``` - -## Development - -Requires a Rust toolchain and Python 3.11+. We use -[uv](https://docs.astral.sh/uv/) and [maturin](https://www.maturin.rs/). - -```bash -# Create a dev environment and install the dev dependencies -uv sync --no-install-package zarrista - -# Build the Rust extension and install it into the environment (debug build) -uv run --no-project maturin develop --uv - -# Or, in release mode: -uv run --no-project maturin develop --uv --release - -# Run the tests -uv run --no-project pytest -``` - -The `--no-project` is annoying but unavoidable in our current setup. Otherwise `uv` will try to build the rust library _in release mode, as a dependency of the project_ before reaching `uv sync` or `uv run`. From 68f0f6280f020b5ecaab5c45c8bcb262c8e7292c Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Mon, 20 Jul 2026 16:27:57 +0200 Subject: [PATCH 6/7] infer intopyobjectref --- src/storage/async.rs | 15 +-------------- 1 file changed, 1 insertion(+), 14 deletions(-) diff --git a/src/storage/async.rs b/src/storage/async.rs index 51e5f01..fae5285 100644 --- a/src/storage/async.rs +++ b/src/storage/async.rs @@ -18,7 +18,7 @@ use zarrs::storage::{ use zarrs_icechunk::AsyncIcechunkStore; use zarrs_object_store::AsyncObjectStore; -#[derive(Clone, IntoPyObject)] +#[derive(Clone, IntoPyObject, IntoPyObjectRef)] pub enum PyAsyncStorage { ObjectStore(PyAsyncObjectStore), Icechunk(PyAsyncIcechunkStore), @@ -54,19 +54,6 @@ impl FromPyObject<'_, '_> for PyAsyncStorage { } } -impl<'py> IntoPyObject<'py> for &PyAsyncStorage { - type Target = PyAny; - type Error = Infallible; - type Output = Bound<'py, Self::Target>; - - fn into_pyobject(self, py: Python<'py>) -> Result { - match self { - PyAsyncStorage::ObjectStore(store) => store.into_pyobject(py), - PyAsyncStorage::Icechunk(store) => store.into_pyobject(py), - } - } -} - impl From for Arc { fn from(s: PyAsyncStorage) -> Self { s.inner() From dca7a34d0418babbc1c4da73bbd28972f12d3ba8 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Mon, 20 Jul 2026 16:48:00 +0200 Subject: [PATCH 7/7] comment --- src/storage/async.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/storage/async.rs b/src/storage/async.rs index fae5285..bc1f902 100644 --- a/src/storage/async.rs +++ b/src/storage/async.rs @@ -63,6 +63,7 @@ impl From for Arc>>, + /// Handle to the original Python object, used for returning the original store to Python pyobj: Arc>, } @@ -101,6 +102,7 @@ impl<'py> IntoPyObject<'py> for &PyAsyncObjectStore { #[derive(Clone)] pub struct PyAsyncIcechunkStore { inner: Arc, + /// Handle to the original Python object, used for returning the original store to Python pyobj: Arc>, }