fix: add StoreConfig super class

Does not actually work, not picked up by stub gen
ilan-gold · Dec 10, 2024 · 0643735 · 0643735
1 parent 71f1698
commit 0643735
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 22 deletions.
diff --git a/python/zarrs/_internal.pyi b/python/zarrs/_internal.pyi
@@ -23,7 +23,11 @@ class CodecPipelineImpl:
         chunk_descriptions: typing.Sequence[
             tuple[
                 tuple[
-                    StoreConfig, str, typing.Sequence[int], str, typing.Sequence[int]
+                    StoreConfigType,
+                    str,
+                    typing.Sequence[int],
+                    str,
+                    typing.Sequence[int],
                 ],
                 typing.Sequence[slice],
                 typing.Sequence[slice],
@@ -34,15 +38,19 @@ class CodecPipelineImpl:
     def retrieve_chunks(
         self,
         chunk_descriptions: typing.Sequence[
-            tuple[StoreConfig, str, typing.Sequence[int], str, typing.Sequence[int]]
+            tuple[StoreConfigType, str, typing.Sequence[int], str, typing.Sequence[int]]
         ],
     ) -> list[numpy.typing.NDArray[numpy.uint8]]: ...
     def store_chunks_with_indices(
         self,
         chunk_descriptions: typing.Sequence[
             tuple[
                 tuple[
-                    StoreConfig, str, typing.Sequence[int], str, typing.Sequence[int]
+                    StoreConfigType,
+                    str,
+                    typing.Sequence[int],
+                    str,
+                    typing.Sequence[int],
                 ],
                 typing.Sequence[slice],
                 typing.Sequence[slice],
@@ -53,7 +61,8 @@ class CodecPipelineImpl:
 
 class FilesystemStoreConfig: ...
 class HttpStoreConfig: ...
+class StoreConfig: ...
 
-class StoreConfig(Enum):
+class StoreConfigType(Enum):
     Filesystem = auto()
     Http = auto()
diff --git a/src/chunk_item.rs b/src/chunk_item.rs
@@ -12,11 +12,11 @@ use zarrs::{
     storage::{MaybeBytes, ReadableWritableListableStorageTraits, StorageError, StoreKey},
 };
 
-use crate::{utils::PyErrExt, StoreConfig};
+use crate::{utils::PyErrExt, StoreConfigType};
 
 pub(crate) type Raw<'a> = (
     // store
-    StoreConfig,
+    StoreConfigType,
     // path
     String,
     // shape
@@ -36,7 +36,7 @@ pub(crate) type RawWithIndices<'a> = (
 );
 
 pub(crate) trait IntoItem<T, S>: std::marker::Sized {
-    fn store_config(&self) -> &StoreConfig;
+    fn store_config(&self) -> &StoreConfigType;
     fn path(&self) -> &str;
     fn into_item(
         self,
@@ -93,7 +93,7 @@ impl ChunksItem for WithSubset {
 }
 
 impl<'a> IntoItem<Basic, ()> for Raw<'a> {
-    fn store_config(&self) -> &StoreConfig {
+    fn store_config(&self) -> &StoreConfigType {
         &self.0
     }
 
@@ -118,7 +118,7 @@ impl<'a> IntoItem<Basic, ()> for Raw<'a> {
 }
 
 impl IntoItem<WithSubset, &[u64]> for RawWithIndices<'_> {
-    fn store_config(&self) -> &StoreConfig {
+    fn store_config(&self) -> &StoreConfigType {
         &self.0 .0
     }
 

diff --git a/src/codec_pipeline_store_filesystem.rs b/src/codec_pipeline_store_filesystem.rs
@@ -4,14 +4,14 @@ use pyo3::{exceptions::PyRuntimeError, pyclass, PyResult};
 use pyo3_stub_gen::derive::gen_stub_pyclass;
 use zarrs::{filesystem::FilesystemStore, storage::ReadableWritableListableStorageTraits};
 
-use crate::{utils::PyErrExt, CodecPipelineStore};
+use crate::{utils::PyErrExt, CodecPipelineStore, StoreConfig};
 
 pub struct CodecPipelineStoreFilesystem {
     store: Arc<FilesystemStore>,
 }
 
 #[gen_stub_pyclass]
-#[pyclass]
+#[pyclass(extends=StoreConfig)]
 pub struct FilesystemStoreConfig {
     root: String,
 }

diff --git a/src/codec_pipeline_store_http.rs b/src/codec_pipeline_store_http.rs
@@ -9,15 +9,15 @@ use zarrs_opendal::AsyncOpendalStore;
 use crate::{
     runtime::{tokio_block_on, TokioBlockOn},
     utils::PyErrExt,
-    CodecPipelineStore,
+    CodecPipelineStore, StoreConfig,
 };
 
 pub struct CodecPipelineStoreHTTP {
     store: Arc<AsyncToSyncStorageAdapter<AsyncOpendalStore, TokioBlockOn>>,
 }
 
 #[gen_stub_pyclass]
-#[pyclass]
+#[pyclass(extends=StoreConfig)]
 pub struct HttpStoreConfig {
     pub root: String,
 }

diff --git a/src/lib.rs b/src/lib.rs
@@ -53,14 +53,18 @@ pub struct CodecPipelineImpl {
     pub(crate) num_threads: usize,
 }
 
+#[gen_stub_pyclass]
+#[pyclass(subclass)]
+pub struct StoreConfig;
+
 #[gen_stub_pyclass_enum]
-enum StoreConfig {
+enum StoreConfigType {
     Filesystem(FilesystemStoreConfig),
     Http(HttpStoreConfig),
     // TODO: Add support for more stores
 }
 
-impl<'py> FromPyObject<'py> for StoreConfig {
+impl<'py> FromPyObject<'py> for StoreConfigType {
     fn extract_bound(store: &Bound<'py, PyAny>) -> PyResult<Self> {
         let name = store.get_type().name()?;
         let name = name.to_str()?;
@@ -70,7 +74,9 @@ impl<'py> FromPyObject<'py> for StoreConfig {
                     .getattr("root")?
                     .call_method("as_posix", (), None)?
                     .extract()?;
-                Ok(StoreConfig::Filesystem(FilesystemStoreConfig::new(root)))
+                Ok(StoreConfigType::Filesystem(FilesystemStoreConfig::new(
+                    root,
+                )))
             }
             "RemoteStore" => {
                 let fs = store.getattr("fs")?;
@@ -80,7 +86,7 @@ impl<'py> FromPyObject<'py> for StoreConfig {
                 let storage_options: HashMap<String, Bound<'py, PyAny>> =
                     fs.getattr("storage_options")?.extract()?;
                 match fs_name {
-                    "HTTPFileSystem" => Ok(StoreConfig::Http(HttpStoreConfig::new(
+                    "HTTPFileSystem" => Ok(StoreConfigType::Http(HttpStoreConfig::new(
                         &path,
                         &storage_options,
                     )?)),
@@ -96,23 +102,23 @@ impl<'py> FromPyObject<'py> for StoreConfig {
     }
 }
 
-impl TryFrom<&StoreConfig> for Arc<dyn CodecPipelineStore> {
+impl TryFrom<&StoreConfigType> for Arc<dyn CodecPipelineStore> {
     type Error = PyErr;
 
-    fn try_from(value: &StoreConfig) -> Result<Self, Self::Error> {
+    fn try_from(value: &StoreConfigType) -> Result<Self, Self::Error> {
         match value {
-            StoreConfig::Filesystem(config) => {
+            StoreConfigType::Filesystem(config) => {
                 Ok(Arc::new(CodecPipelineStoreFilesystem::new(config)?))
             }
-            StoreConfig::Http(config) => Ok(Arc::new(CodecPipelineStoreHTTP::new(config)?)),
+            StoreConfigType::Http(config) => Ok(Arc::new(CodecPipelineStoreHTTP::new(config)?)),
         }
     }
 }
 
 impl CodecPipelineImpl {
     fn get_store_from_config(
         &self,
-        config: &StoreConfig,
+        config: &StoreConfigType,
     ) -> PyResult<Arc<dyn ReadableWritableListableStorageTraits>> {
         let mut gstore = self.store.lock().map_err(|_| {
             PyErr::new::<PyRuntimeError, _>("failed to lock the store mutex".to_string())