Skip to content

feat: Add Store.replace_prefix - #762

Draft
kylebarron wants to merge 5 commits into
mainfrom
kyle/replace-prefix
Draft

feat: Add Store.replace_prefix#762
kylebarron wants to merge 5 commits into
mainfrom
kyle/replace-prefix

Conversation

@kylebarron

@kylebarron kylebarron commented Aug 5, 2026

Copy link
Copy Markdown
Member

Change list

  • Adds Store.replace_prefix for each store. This swaps out the prefix on the store, returning a new store instance but, crucially, it shares the same connection pool as the previous store instance. So this is very cheap to call.

This doesn't currently compile because it requires apache/arrow-rs-object-store#825; once that merges we can pin to a git tag upstream.

cc @d-v-b

@ds-release-bot ds-release-bot Bot added the feat label Aug 5, 2026
@kylebarron

kylebarron commented Aug 21, 2026

Copy link
Copy Markdown
Member Author

Comment written by Claude


replace_prefix doesn't work with our own store classes

Every public store class is a Python subclass of the Rust pyclass:

class S3Store(ObjectStoreMethods, _store.S3Store): ...

and get / put / list / copy / head (and all the _async variants) live on the
ObjectStoreMethods mixin — the pyclass itself has none of them.

replace_prefix constructs the Rust Self directly, so PyO3 always returns the base
_store.S3Store. The returned object therefore has no store API at all:

>>> s = obstore.store.S3Store("bucket", prefix="a")
>>> n = s.replace_prefix("b")
>>> type(n)
<class 'obstore._store.S3Store'>   # not obstore.store.S3Store
>>> n.get("path.txt")
AttributeError: 'builtins.S3Store' object has no attribute 'get'

This affects every caller, not just user-defined subclasses. The .pyi stubs annotate
-> Self, so nothing catches it at type-check time either. from_url already dodges this by
deliberately passing back through Python (cls.call(...)); replace_prefix doesn't.

Why it isn't a one-line fix

Routing through the constructor the way from_url does would rebuild the store, and
AmazonS3Builder::build() creates a fresh reqwest::Client — a new connection pool. That
defeats the entire point of this API. (Injecting a shared HttpConnector doesn't help either:
it's a factory, so each build() still mints a client.)

Sharing the pool requires cloning the existing AmazonS3 in Rust; returning a subclass requires
going through that subclass's __new__. PyO3 has no public API to allocate a
dynamically-known subtype — create_class_object_of_type is pub(crate), and #[new] is the
only place PyO3 hands you the subtype.

Approach that gets both

A private _clone_from keyword on __new__ carrying an already-built store:

if let Some(source) = _clone_from {
    return Ok(Self {
        store: Arc::new(source.store.replace_prefix(prefix.clone())),
        config: source.config.replace_prefix(prefix),
    });
}

and replace_prefix dispatching through cls.__new__(cls, ...) rather than cls(...), so the
subclass's __init__ isn't handed arguments it never agreed to accept. That's the same path
pickle already takes here via __getnewargs_ex__.

LocalStore uses the same __new__ dispatch but no _clone_from — no pool to preserve, and
mkdir still needs to run.

Tests

None currently on this branch. The patch below adds:

  • tests/store/test_replace_prefix.py — 31 tests, parametrized across S3/Azure/GCS plus
    LocalStore: prefix replaced (not appended), None clears it, config inherited, pickle
    round-trip, subclass preserved with __init__ not re-run, and an end-to-end minio test
    writing through both stores and reading the keys back via an unprefixed store.
  • Three Rust unit tests in prefix.rs. The load-bearing one is
    replace_prefix_shares_the_underlying_store: it clones while the store is empty, writes
    through the original, and reads back through the replacement's inner store — which only
    succeeds if the inner store is genuinely shared.

455 passed, 2 xfailed (pytest) · 3 passed (cargo test) · pyright, ruff, clippy -D warnings
and cargo fmt --check all clean.


Patch — applies cleanly on top of this branch

Save the block below as replace-prefix-subclass.patch and:

$ git apply replace-prefix-subclass.patch
$ uv run maturin dev -m obstore/Cargo.toml
$ uv run pytest tests/store/test_replace_prefix.py
diff --git a/obstore/python/obstore/_store/__init__.pyi b/obstore/python/obstore/_store/__init__.pyi
index 620a62f..d3d4aa7 100644
--- a/obstore/python/obstore/_store/__init__.pyi
+++ b/obstore/python/obstore/_store/__init__.pyi
@@ -193,6 +193,10 @@ class LocalStore:
         equivalent to constructing a new [`LocalStore`][obstore.store.LocalStore]
         directly.
 
+        If this store is an instance of a subclass, the new store is an instance of that
+        same subclass; note that its `__init__` is not called, in the same way that
+        unpickling a store does not call `__init__`.
+
         The new prefix fully replaces the existing one; it is not appended to it.
 
         **Example:**
diff --git a/obstore/python/obstore/_store/_aws.pyi b/obstore/python/obstore/_store/_aws.pyi
index 32eea19..9e43ba3 100644
--- a/obstore/python/obstore/_store/_aws.pyi
+++ b/obstore/python/obstore/_store/_aws.pyi
@@ -628,7 +628,10 @@ class S3Store:
         connections stay warm. This makes it cheap to create many stores pointing at
         different prefixes of the same bucket.
 
-        All other configuration is inherited from this store.
+        All other configuration is inherited from this store. If this store is an instance
+        of a subclass, the new store is an instance of that same subclass; note that its
+        `__init__` is not called, in the same way that unpickling a store does not call
+        `__init__`.
 
         The new prefix fully replaces the existing one; it is not appended to it. It is
         always interpreted relative to the root of the bucket.
diff --git a/obstore/python/obstore/_store/_azure.pyi b/obstore/python/obstore/_store/_azure.pyi
index 46f00e6..eb8f3b4 100644
--- a/obstore/python/obstore/_store/_azure.pyi
+++ b/obstore/python/obstore/_store/_azure.pyi
@@ -456,7 +456,10 @@ class AzureStore:
         connections stay warm. This makes it cheap to create many stores pointing at
         different prefixes of the same container.
 
-        All other configuration is inherited from this store.
+        All other configuration is inherited from this store. If this store is an instance
+        of a subclass, the new store is an instance of that same subclass; note that its
+        `__init__` is not called, in the same way that unpickling a store does not call
+        `__init__`.
 
         The new prefix fully replaces the existing one; it is not appended to it. It is
         always interpreted relative to the root of the container.
diff --git a/obstore/python/obstore/_store/_gcs.pyi b/obstore/python/obstore/_store/_gcs.pyi
index fa7a2b9..c549299 100644
--- a/obstore/python/obstore/_store/_gcs.pyi
+++ b/obstore/python/obstore/_store/_gcs.pyi
@@ -234,7 +234,10 @@ class GCSStore:
         connections stay warm. This makes it cheap to create many stores pointing at
         different prefixes of the same bucket.
 
-        All other configuration is inherited from this store.
+        All other configuration is inherited from this store. If this store is an instance
+        of a subclass, the new store is an instance of that same subclass; note that its
+        `__init__` is not called, in the same way that unpickling a store does not call
+        `__init__`.
 
         The new prefix fully replaces the existing one; it is not appended to it. It is
         always interpreted relative to the root of the bucket.
diff --git a/pyo3-object_store/src/aws/store.rs b/pyo3-object_store/src/aws/store.rs
index ef949ef..4b95f7e 100644
--- a/pyo3-object_store/src/aws/store.rs
+++ b/pyo3-object_store/src/aws/store.rs
@@ -95,7 +95,10 @@ impl PyS3Store {
 impl PyS3Store {
     // Create from parameters
     #[new]
-    #[pyo3(signature = (bucket=None, *, prefix=None, config=None, client_options=None, retry_config=None, credential_provider=None, **kwargs))]
+    #[pyo3(signature = (bucket=None, *, prefix=None, config=None, client_options=None, retry_config=None, credential_provider=None, _clone_from=None, **kwargs))]
+    // `_clone_from` pushes this one over clippy's limit; splitting the constructor isn't an
+    // option because this signature is the public Python API.
+    #[allow(clippy::too_many_arguments)]
     fn new(
         bucket: Option<String>,
         prefix: Option<PyPath>,
@@ -103,8 +106,20 @@ impl PyS3Store {
         client_options: Option<PyClientOptions>,
         retry_config: Option<PyRetryConfig>,
         credential_provider: Option<PyAWSCredentialProvider>,
+        // Private, and set only by `replace_prefix`. PyO3 can allocate a Python
+        // subclass only through `__new__`, so an already-built store has to arrive here
+        // as an argument; reusing it is what keeps the HTTP client, and therefore its
+        // connection pool, shared. Every other argument is ignored when it's set.
+        _clone_from: Option<PyRef<'_, Self>>,
         kwargs: Option<PyAmazonS3Config>,
     ) -> PyObjectStoreResult<Self> {
+        if let Some(source) = _clone_from {
+            return Ok(Self {
+                store: Arc::new(source.store.replace_prefix(prefix.clone())),
+                config: source.config.replace_prefix(prefix),
+            });
+        }
+
         let mut builder = AmazonS3Builder::from_env();
         let mut config = config.unwrap_or_default();
 
@@ -226,11 +241,22 @@ impl PyS3Store {
         self.config.credential_provider.as_ref()
     }
 
-    fn replace_prefix(&self, prefix: Option<PyPath>) -> PyObjectStoreResult<Self> {
-        Ok(Self {
-            store: Arc::new(self.store.replace_prefix(prefix.clone())),
-            config: self.config.replace_prefix(prefix),
-        })
+    fn replace_prefix<'py>(
+        slf: &Bound<'py, Self>,
+        prefix: Option<PyPath>,
+    ) -> PyObjectStoreResult<Bound<'py, PyAny>> {
+        let py = slf.py();
+        let cls = slf.get_type();
+
+        let kwargs = PyDict::new(py);
+        kwargs.set_item(intern!(py, "prefix"), prefix)?;
+        kwargs.set_item(intern!(py, "_clone_from"), slf)?;
+
+        // Note: we pass **back** through Python so that if `cls` is a subclass, we instantiate
+        // the subclass. We call `__new__` directly rather than `cls(...)` so that we don't invoke
+        // a subclass's `__init__` with arguments it doesn't expect; this is the same path pickle
+        // takes via `__getnewargs_ex__`.
+        Ok(cls.call_method(intern!(py, "__new__"), (&cls,), Some(&kwargs))?)
     }
 
     #[getter]
diff --git a/pyo3-object_store/src/azure/store.rs b/pyo3-object_store/src/azure/store.rs
index 0df27a4..8c43068 100644
--- a/pyo3-object_store/src/azure/store.rs
+++ b/pyo3-object_store/src/azure/store.rs
@@ -101,7 +101,10 @@ impl PyAzureStore {
 impl PyAzureStore {
     // Create from parameters
     #[new]
-    #[pyo3(signature = (container_name=None, *, prefix=None, config=None, client_options=None, retry_config=None, credential_provider=None, **kwargs))]
+    #[pyo3(signature = (container_name=None, *, prefix=None, config=None, client_options=None, retry_config=None, credential_provider=None, _clone_from=None, **kwargs))]
+    // `_clone_from` pushes this one over clippy's limit; splitting the constructor isn't an
+    // option because this signature is the public Python API.
+    #[allow(clippy::too_many_arguments)]
     fn new(
         container_name: Option<String>,
         mut prefix: Option<PyPath>,
@@ -109,8 +112,20 @@ impl PyAzureStore {
         client_options: Option<PyClientOptions>,
         retry_config: Option<PyRetryConfig>,
         credential_provider: Option<PyAzureCredentialProvider>,
+        // Private, and set only by `replace_prefix`. PyO3 can allocate a Python
+        // subclass only through `__new__`, so an already-built store has to arrive here
+        // as an argument; reusing it is what keeps the HTTP client, and therefore its
+        // connection pool, shared. Every other argument is ignored when it's set.
+        _clone_from: Option<PyRef<'_, Self>>,
         kwargs: Option<PyAzureConfig>,
     ) -> PyObjectStoreResult<Self> {
+        if let Some(source) = _clone_from {
+            return Ok(Self {
+                store: Arc::new(source.store.replace_prefix(prefix.clone())),
+                config: source.config.replace_prefix(prefix),
+            });
+        }
+
         let mut builder = MicrosoftAzureBuilder::from_env();
         let mut config = config.unwrap_or_default();
 
@@ -248,11 +263,22 @@ impl PyAzureStore {
         self.config.credential_provider.as_ref()
     }
 
-    fn replace_prefix(&self, prefix: Option<PyPath>) -> PyObjectStoreResult<Self> {
-        Ok(Self {
-            store: Arc::new(self.store.replace_prefix(prefix.clone())),
-            config: self.config.replace_prefix(prefix),
-        })
+    fn replace_prefix<'py>(
+        slf: &Bound<'py, Self>,
+        prefix: Option<PyPath>,
+    ) -> PyObjectStoreResult<Bound<'py, PyAny>> {
+        let py = slf.py();
+        let cls = slf.get_type();
+
+        let kwargs = PyDict::new(py);
+        kwargs.set_item(intern!(py, "prefix"), prefix)?;
+        kwargs.set_item(intern!(py, "_clone_from"), slf)?;
+
+        // Note: we pass **back** through Python so that if `cls` is a subclass, we instantiate
+        // the subclass. We call `__new__` directly rather than `cls(...)` so that we don't invoke
+        // a subclass's `__init__` with arguments it doesn't expect; this is the same path pickle
+        // takes via `__getnewargs_ex__`.
+        Ok(cls.call_method(intern!(py, "__new__"), (&cls,), Some(&kwargs))?)
     }
 
     #[getter]
diff --git a/pyo3-object_store/src/gcp/store.rs b/pyo3-object_store/src/gcp/store.rs
index 99a38a1..71fdf4b 100644
--- a/pyo3-object_store/src/gcp/store.rs
+++ b/pyo3-object_store/src/gcp/store.rs
@@ -93,7 +93,10 @@ impl PyGCSStore {
 impl PyGCSStore {
     // Create from parameters
     #[new]
-    #[pyo3(signature = (bucket=None, *, prefix=None, config=None, client_options=None, retry_config=None, credential_provider=None, **kwargs))]
+    #[pyo3(signature = (bucket=None, *, prefix=None, config=None, client_options=None, retry_config=None, credential_provider=None, _clone_from=None, **kwargs))]
+    // `_clone_from` pushes this one over clippy's limit; splitting the constructor isn't an
+    // option because this signature is the public Python API.
+    #[allow(clippy::too_many_arguments)]
     fn new(
         bucket: Option<String>,
         prefix: Option<PyPath>,
@@ -101,8 +104,20 @@ impl PyGCSStore {
         client_options: Option<PyClientOptions>,
         retry_config: Option<PyRetryConfig>,
         credential_provider: Option<PyGcpCredentialProvider>,
+        // Private, and set only by `replace_prefix`. PyO3 can allocate a Python
+        // subclass only through `__new__`, so an already-built store has to arrive here
+        // as an argument; reusing it is what keeps the HTTP client, and therefore its
+        // connection pool, shared. Every other argument is ignored when it's set.
+        _clone_from: Option<PyRef<'_, Self>>,
         kwargs: Option<PyGoogleConfig>,
     ) -> PyObjectStoreResult<Self> {
+        if let Some(source) = _clone_from {
+            return Ok(Self {
+                store: Arc::new(source.store.replace_prefix(prefix.clone())),
+                config: source.config.replace_prefix(prefix),
+            });
+        }
+
         let mut builder = GoogleCloudStorageBuilder::from_env();
         let mut config = config.unwrap_or_default();
         if let Some(bucket) = bucket.clone() {
@@ -212,11 +227,22 @@ impl PyGCSStore {
         self.config.credential_provider.as_ref()
     }
 
-    fn replace_prefix(&self, prefix: Option<PyPath>) -> PyObjectStoreResult<Self> {
-        Ok(Self {
-            store: Arc::new(self.store.replace_prefix(prefix.clone())),
-            config: self.config.replace_prefix(prefix),
-        })
+    fn replace_prefix<'py>(
+        slf: &Bound<'py, Self>,
+        prefix: Option<PyPath>,
+    ) -> PyObjectStoreResult<Bound<'py, PyAny>> {
+        let py = slf.py();
+        let cls = slf.get_type();
+
+        let kwargs = PyDict::new(py);
+        kwargs.set_item(intern!(py, "prefix"), prefix)?;
+        kwargs.set_item(intern!(py, "_clone_from"), slf)?;
+
+        // Note: we pass **back** through Python so that if `cls` is a subclass, we instantiate
+        // the subclass. We call `__new__` directly rather than `cls(...)` so that we don't invoke
+        // a subclass's `__init__` with arguments it doesn't expect; this is the same path pickle
+        // takes via `__getnewargs_ex__`.
+        Ok(cls.call_method(intern!(py, "__new__"), (&cls,), Some(&kwargs))?)
     }
 
     #[getter]
diff --git a/pyo3-object_store/src/local.rs b/pyo3-object_store/src/local.rs
index 61d40c7..8719eed 100644
--- a/pyo3-object_store/src/local.rs
+++ b/pyo3-object_store/src/local.rs
@@ -141,10 +141,22 @@ impl PyLocalStore {
         }
     }
 
-    fn replace_prefix(&self, prefix: Option<std::path::PathBuf>) -> PyObjectStoreResult<Self> {
-        // Here we use Self::new instead of `replace_prefix` as on the other stores because 1) this
-        // doesn't use a MaybePrefixedStore wrapper and 2) there's no underlying connection pool we
-        // need to reuse.
-        Self::new(prefix, self.config.automatic_cleanup, self.config.mkdir)
+    fn replace_prefix<'py>(
+        slf: &Bound<'py, Self>,
+        prefix: Option<std::path::PathBuf>,
+    ) -> PyObjectStoreResult<Bound<'py, PyAny>> {
+        let py = slf.py();
+        let config = &slf.get().config;
+
+        // Note: we pass **back** through Python so that if this is a subclass, we instantiate the
+        // subclass, calling `__new__` directly so that we don't invoke a subclass's `__init__`
+        // with arguments it doesn't expect. Unlike the remote stores there's no `_clone_from`
+        // path, because there's no underlying connection pool we'd lose by constructing the store
+        // from scratch.
+        let cls = slf.get_type();
+        let kwargs = PyDict::new(py);
+        kwargs.set_item(intern!(py, "automatic_cleanup"), config.automatic_cleanup)?;
+        kwargs.set_item(intern!(py, "mkdir"), config.mkdir)?;
+        Ok(cls.call_method(intern!(py, "__new__"), (&cls, prefix), Some(&kwargs))?)
     }
 }
diff --git a/pyo3-object_store/src/prefix.rs b/pyo3-object_store/src/prefix.rs
index 2e6739c..3bf109c 100644
--- a/pyo3-object_store/src/prefix.rs
+++ b/pyo3-object_store/src/prefix.rs
@@ -269,3 +269,70 @@ impl<T: ObjectStore + Signer> Signer for MaybePrefixedStore<T> {
         })
     }
 }
+
+#[cfg(test)]
+mod test {
+    use object_store::memory::InMemory;
+    use object_store::ObjectStoreExt;
+
+    use super::*;
+
+    /// `replace_prefix` must share the underlying store rather than construct a new one; for the
+    /// remote stores that sharing is what keeps the HTTP client, and therefore its connection
+    /// pool, alive. Here we observe it through `InMemory`, whose clones share backing storage.
+    #[test]
+    fn replace_prefix_shares_the_underlying_store() {
+        let store = MaybePrefixedStore::new(InMemory::new(), Some(Path::from("a")));
+        let replaced = store.replace_prefix(Some(Path::from("b")));
+
+        tokio::runtime::Runtime::new().unwrap().block_on(async {
+            store
+                .put(&Path::from("x"), PutPayload::from_static(b"hello"))
+                .await
+                .unwrap();
+
+            // The write landed at `a/x`. `replaced` can read it back through its own inner store
+            // only because that inner store is the *same* `InMemory`, not a newly built one.
+            let bytes = replaced
+                .inner()
+                .get(&Path::from("a/x"))
+                .await
+                .unwrap()
+                .bytes()
+                .await
+                .unwrap();
+            assert_eq!(bytes.as_ref(), b"hello");
+        });
+    }
+
+    #[test]
+    fn replace_prefix_replaces_rather_than_appends() {
+        let store = MaybePrefixedStore::new(InMemory::new(), Some(Path::from("a")));
+        let replaced = store.replace_prefix(Some(Path::from("b")));
+
+        tokio::runtime::Runtime::new().unwrap().block_on(async {
+            replaced
+                .put(&Path::from("x"), PutPayload::from_static(b"hello"))
+                .await
+                .unwrap();
+
+            assert!(replaced.inner().get(&Path::from("b/x")).await.is_ok());
+            assert!(replaced.inner().get(&Path::from("a/b/x")).await.is_err());
+        });
+    }
+
+    #[test]
+    fn replace_prefix_with_none_clears_the_prefix() {
+        let store = MaybePrefixedStore::new(InMemory::new(), Some(Path::from("a")));
+        let replaced = store.replace_prefix(None::<Path>);
+
+        tokio::runtime::Runtime::new().unwrap().block_on(async {
+            replaced
+                .put(&Path::from("x"), PutPayload::from_static(b"hello"))
+                .await
+                .unwrap();
+
+            assert!(replaced.inner().get(&Path::from("x")).await.is_ok());
+        });
+    }
+}
diff --git a/tests/store/test_replace_prefix.py b/tests/store/test_replace_prefix.py
new file mode 100644
index 0000000..60a3dda
--- /dev/null
+++ b/tests/store/test_replace_prefix.py
@@ -0,0 +1,199 @@
+"""Tests for `Store.replace_prefix` across each store backend."""
+
+from __future__ import annotations
+
+import pickle
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import pytest
+
+from obstore.store import AzureStore, GCSStore, LocalStore, S3Store
+
+if TYPE_CHECKING:
+    from obstore.store import ClientConfig, S3Config
+
+
+def remote_stores() -> list[S3Store | AzureStore | GCSStore]:
+    """One store per remote backend, each configured with a prefix."""
+    return [
+        S3Store(
+            "bucket",
+            prefix="data/2024",
+            region="us-east-1",
+            skip_signature=True,
+            client_options={"timeout": "10s"},
+            retry_config={"max_retries": 5},
+        ),
+        AzureStore(
+            "container",
+            prefix="data/2024",
+            account_name="account",
+            skip_signature=True,
+            client_options={"timeout": "10s"},
+            retry_config={"max_retries": 5},
+        ),
+        GCSStore(
+            "bucket",
+            prefix="data/2024",
+            skip_signature=True,
+            client_options={"timeout": "10s"},
+            retry_config={"max_retries": 5},
+        ),
+    ]
+
+
+@pytest.fixture(params=remote_stores(), ids=lambda store: type(store).__name__)
+def remote_store(request: pytest.FixtureRequest) -> S3Store | AzureStore | GCSStore:
+    return request.param
+
+
+def test_replaces_prefix(remote_store: S3Store | AzureStore | GCSStore):
+    assert remote_store.replace_prefix("data/2025").prefix == "data/2025"
+
+
+def test_replaces_prefix_rather_than_appending(
+    remote_store: S3Store | AzureStore | GCSStore,
+):
+    twice = remote_store.replace_prefix("data/2025").replace_prefix("data/2026")
+    assert twice.prefix == "data/2026"
+
+
+def test_none_clears_prefix(remote_store: S3Store | AzureStore | GCSStore):
+    assert remote_store.replace_prefix(None).prefix is None
+
+
+def test_original_store_is_unchanged(remote_store: S3Store | AzureStore | GCSStore):
+    remote_store.replace_prefix("data/2025")
+    assert remote_store.prefix == "data/2024"
+
+
+def test_other_config_is_inherited(remote_store: S3Store | AzureStore | GCSStore):
+    new_store = remote_store.replace_prefix("data/2025")
+    assert new_store.config == remote_store.config
+    assert new_store.client_options == remote_store.client_options
+    assert new_store.retry_config == remote_store.retry_config
+
+
+def test_eq_matches_a_directly_constructed_store(
+    remote_store: S3Store | AzureStore | GCSStore,
+):
+    new_store = remote_store.replace_prefix("data/2025")
+    directly = type(remote_store)(
+        prefix="data/2025",
+        config=remote_store.config,  # type: ignore[arg-type]
+        client_options=remote_store.client_options,
+        retry_config=remote_store.retry_config,
+    )
+    assert new_store == directly
+
+
+def test_pickle_round_trip(remote_store: S3Store | AzureStore | GCSStore):
+    """The pickling config must stay in sync with the underlying store's prefix."""
+    new_store = remote_store.replace_prefix("data/2025")
+    restored = pickle.loads(pickle.dumps(new_store))
+    assert restored.prefix == "data/2025"
+    assert restored == new_store
+
+
+def test_preserves_subclass(remote_store: S3Store | AzureStore | GCSStore):
+    init_calls = []
+
+    class Subclass(type(remote_store)):  # type: ignore[misc]
+        def __init__(self, *_args: object, **_kwargs: object) -> None:
+            init_calls.append(1)
+
+    store = Subclass(
+        prefix="data/2024",
+        config=remote_store.config,  # type: ignore[arg-type]
+        client_options=remote_store.client_options,
+        retry_config=remote_store.retry_config,
+    )
+    assert len(init_calls) == 1
+
+    new_store = store.replace_prefix("data/2025")
+    assert type(new_store) is Subclass
+    assert new_store.prefix == "data/2025"
+    assert new_store.config == store.config
+    # `__init__` is not re-run, just as unpickling a store does not call it.
+    assert len(init_calls) == 1
+
+
+
+def test_local_replaces_prefix(tmp_path: Path):
+    (tmp_path / "2024").mkdir()
+    (tmp_path / "2025").mkdir()
+
+    store = LocalStore(tmp_path / "2024", automatic_cleanup=True)
+    new_store = store.replace_prefix(tmp_path / "2025")
+
+    assert new_store.prefix == tmp_path / "2025"
+    assert isinstance(new_store.prefix, Path)
+    # The original is untouched
+    assert store.prefix == tmp_path / "2024"
+    # And the rest of the config is inherited
+    assert new_store == LocalStore(tmp_path / "2025", automatic_cleanup=True)
+
+
+def test_local_none_clears_prefix(tmp_path: Path):
+    assert LocalStore(tmp_path).replace_prefix(None).prefix is None
+
+
+def test_local_mkdir_is_inherited(tmp_path: Path):
+    store = LocalStore(tmp_path / "2024", mkdir=True)
+    new_dir = tmp_path / "2025"
+    assert not new_dir.exists()
+
+    store.replace_prefix(new_dir)
+    assert new_dir.exists()
+
+
+def test_local_writes_to_the_new_prefix(tmp_path: Path):
+    store = LocalStore(tmp_path / "2024", mkdir=True)
+    new_store = store.replace_prefix(tmp_path / "2025")
+    new_store.put("afile.txt", b"hello world")
+
+    assert (tmp_path / "2025" / "afile.txt").read_bytes() == b"hello world"
+    assert not (tmp_path / "2024" / "afile.txt").exists()
+
+
+def test_local_pickle_round_trip(tmp_path: Path):
+    store = LocalStore(tmp_path / "2024", mkdir=True)
+    new_store = store.replace_prefix(tmp_path / "2025")
+    restored: LocalStore = pickle.loads(pickle.dumps(new_store))
+    assert restored.prefix == tmp_path / "2025"
+    assert restored == new_store
+
+
+def test_local_preserves_subclass(tmp_path: Path):
+    init_calls = []
+
+    class MyLocalStore(LocalStore):
+        def __init__(self, *_args: object, **_kwargs: object) -> None:
+            init_calls.append(1)
+
+    store = MyLocalStore(tmp_path / "2024", mkdir=True)
+    assert len(init_calls) == 1
+
+    new_store = store.replace_prefix(tmp_path / "2025")
+    assert type(new_store) is MyLocalStore
+    assert new_store.prefix == tmp_path / "2025"
+    # `__init__` is not re-run, just as unpickling a store does not call it.
+    assert len(init_calls) == 1
+
+
+def test_writes_to_the_new_prefix(minio_bucket: tuple[S3Config, ClientConfig]):
+    """End-to-end: the underlying store, not just the config, is re-prefixed."""
+    config, client_options = minio_bucket
+    store = S3Store(prefix="data/2024", config=config, client_options=client_options)
+    new_store = store.replace_prefix("data/2025")
+
+    store.put("afile.txt", b"2024")
+    new_store.put("afile.txt", b"2025")
+
+    unprefixed = S3Store(config=config, client_options=client_options)
+    assert unprefixed.get("data/2024/afile.txt").bytes() == b"2024"
+    assert unprefixed.get("data/2025/afile.txt").bytes() == b"2025"
+
+    # Listing through the new store only sees the new prefix
+    assert [obj["path"] for obj in new_store.list().collect()] == ["afile.txt"]

@kylebarron
kylebarron marked this pull request as draft August 21, 2026 21:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant