From da1a56536a606dda32c5d56ee02178cf8d28d848 Mon Sep 17 00:00:00 2001
From: Kyle Barron <kyle@developmentseed.org>
Date: Wed, 5 Feb 2025 15:07:12 -0500
Subject: [PATCH] Fix pickle for local and http store (#223)

---
 Cargo.lock                     | 2 +-
 obstore/Cargo.toml             | 2 +-
 pyo3-object_store/src/http.rs  | 2 +-
 pyo3-object_store/src/local.rs | 2 +-
 tests/store/test_http.py       | 9 +++++++++
 tests/store/test_local.py      | 9 +++++++++
 uv.lock                        | 2 +-
 7 files changed, 23 insertions(+), 5 deletions(-)
 create mode 100644 tests/store/test_http.py

diff --git a/Cargo.lock b/Cargo.lock
index 968ae5e2..c5cfd22b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1865,7 +1865,7 @@ dependencies = [
 
 [[package]]
 name = "obstore"
-version = "0.4.0-beta.3"
+version = "0.4.0-beta.4"
 dependencies = [
  "arrow",
  "bytes",
diff --git a/obstore/Cargo.toml b/obstore/Cargo.toml
index 9208ef76..b5d59c5a 100644
--- a/obstore/Cargo.toml
+++ b/obstore/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "obstore"
-version = "0.4.0-beta.3"
+version = "0.4.0-beta.4"
 authors = { workspace = true }
 edition = { workspace = true }
 description = "The simplest, highest-throughput interface to Amazon S3, Google Cloud Storage, Azure Blob Storage, and S3-compliant APIs like Cloudflare R2."
diff --git a/pyo3-object_store/src/http.rs b/pyo3-object_store/src/http.rs
index 251a4515..9dcc892b 100644
--- a/pyo3-object_store/src/http.rs
+++ b/pyo3-object_store/src/http.rs
@@ -32,7 +32,7 @@ impl HTTPConfig {
 }
 
 /// A Python-facing wrapper around a [`HttpStore`].
-#[pyclass(name = "HTTPStore", frozen)]
+#[pyclass(name = "HTTPStore", module = "obstore.store", frozen)]
 pub struct PyHttpStore {
     // Note: we don't need to wrap this in a MaybePrefixedStore because the HttpStore manages its
     // own prefix.
diff --git a/pyo3-object_store/src/local.rs b/pyo3-object_store/src/local.rs
index 26dcf39e..11f457d9 100644
--- a/pyo3-object_store/src/local.rs
+++ b/pyo3-object_store/src/local.rs
@@ -30,7 +30,7 @@ impl LocalConfig {
 }
 
 /// A Python-facing wrapper around a [`LocalFileSystem`].
-#[pyclass(name = "LocalStore", frozen)]
+#[pyclass(name = "LocalStore", module = "obstore.store", frozen)]
 pub struct PyLocalStore {
     store: Arc<LocalFileSystem>,
     config: LocalConfig,
diff --git a/tests/store/test_http.py b/tests/store/test_http.py
new file mode 100644
index 00000000..e72cb553
--- /dev/null
+++ b/tests/store/test_http.py
@@ -0,0 +1,9 @@
+import pickle
+
+from obstore.store import HTTPStore
+
+
+def test_pickle():
+    store = HTTPStore.from_url("https://example.com")
+    new_store: HTTPStore = pickle.loads(pickle.dumps(store))
+    assert store.url == new_store.url
diff --git a/tests/store/test_local.py b/tests/store/test_local.py
index 90f84cab..0219e573 100644
--- a/tests/store/test_local.py
+++ b/tests/store/test_local.py
@@ -1,3 +1,4 @@
+import pickle
 import tempfile
 from pathlib import Path
 
@@ -63,3 +64,11 @@ def test_prefix_property():
     assert isinstance(store.prefix, Path)
     # Can pass it back to the store init
     LocalStore(store.prefix)
+
+
+def test_pickle():
+    tmpdir = Path(tempfile.gettempdir())
+    store = LocalStore(tmpdir)
+    obs.put(store, "path.txt", b"foo")
+    new_store: LocalStore = pickle.loads(pickle.dumps(store))
+    assert obs.get(new_store, "path.txt").bytes() == b"foo"
diff --git a/uv.lock b/uv.lock
index 1881a0c3..2264b86f 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1350,7 +1350,6 @@ wheels = [
 
 [[package]]
 name = "obstore"
-version = "0.3.0b10"
 source = { directory = "obstore" }
 
 [[package]]
@@ -2458,6 +2457,7 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/b5/e8/dbf020b4d98251a9860752a094d09a65e1b436ad181faf929983f697048f/watchdog-6.0.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:20ffe5b202af80ab4266dcd3e91aae72bf2da48c0d33bdb15c66658e685e94e2", size = 79078 },
     { url = "https://files.pythonhosted.org/packages/07/f6/d0e5b343768e8bcb4cda79f0f2f55051bf26177ecd5651f84c07567461cf/watchdog-6.0.0-py3-none-win32.whl", hash = "sha256:07df1fdd701c5d4c8e55ef6cf55b8f0120fe1aef7ef39a1c6fc6bc2e606d517a", size = 79065 },
     { url = "https://files.pythonhosted.org/packages/db/d9/c495884c6e548fce18a8f40568ff120bc3a4b7b99813081c8ac0c936fa64/watchdog-6.0.0-py3-none-win_amd64.whl", hash = "sha256:cbafb470cf848d93b5d013e2ecb245d4aa1c8fd0504e863ccefa32445359d680", size = 79070 },
+    { url = "https://files.pythonhosted.org/packages/33/e8/e40370e6d74ddba47f002a32919d91310d6074130fe4e17dabcafc15cbf1/watchdog-6.0.0-py3-none-win_ia64.whl", hash = "sha256:a1914259fa9e1454315171103c6a30961236f508b9b623eae470268bbcc6a22f", size = 79067 },
 ]
 
 [[package]]