Set store to read only after unpickling (#405)

* Set store to read only after unpickling

Closes #383
xref #185

* tpying

icechunk-python/python/icechunk/__init__.py

@@ -1,5 +1,6 @@
 # module
-from collections.abc import AsyncGenerator, AsyncIterator, Iterable
+import contextlib
+from collections.abc import AsyncGenerator, AsyncIterator, Generator, Iterable
 from typing import Any, Self
 
 from icechunk._icechunk_python import (
@@ -33,6 +34,7 @@
 
 class IcechunkStore(Store, SyncMixin):
     _store: PyIcechunkStore
+    _pickle_preserves_read_only: bool
 
     @classmethod
     async def open(cls, *args: Any, **kwargs: Any) -> Self:
@@ -89,6 +91,7 @@ def __init__(
                 "An IcechunkStore should not be created with the default constructor, instead use either the create or open_existing class methods."
             )
         self._store = store
+        self._pickle_preserves_read_only = False
 
     @classmethod
     def open_existing(
@@ -152,6 +155,8 @@ def __getstate__(self) -> object:
         # we serialize the Rust store as bytes
         d = self.__dict__.copy()
         d["_store"] = self._store.as_bytes()
+        if not self._pickle_preserves_read_only:
+            d["_read_only"] = True
         return d
 
     def __setstate__(self, state: Any) -> None:
@@ -161,6 +166,18 @@ def __setstate__(self, state: Any) -> None:
         state["_store"] = pyicechunk_store_from_bytes(store_repr, read_only)
         self.__dict__ = state
 
+    @contextlib.contextmanager
+    def preserve_read_only(self) -> Generator[None, None, None]:
+        """
+        Context manager to allow unpickling this store preserving `read_only` status.
+        By default, stores are set to read-only after unpickling.
+        """
+        try:
+            self._pickle_preserves_read_only = True
+            yield
+        finally:
+            self._pickle_preserves_read_only = False
+
     def as_read_only(self) -> Self:
         """Return a read-only version of this store."""
         new_store = self._store.with_read_only(read_only=True)

icechunk-python/tests/test_pickle.py

@@ -20,6 +20,19 @@ async def tmp_store(tmpdir):
     store.close()
 
 
+async def test_pickle_read_only(tmp_store):
+    assert tmp_store._read_only is False
+
+    roundtripped = pickle.loads(pickle.dumps(tmp_store))
+    assert roundtripped._read_only is True
+
+    with tmp_store.preserve_read_only():
+        roundtripped = pickle.loads(pickle.dumps(tmp_store))
+        assert roundtripped._read_only is False
+
+    assert tmp_store._read_only is False
+
+
 async def test_pickle(tmp_store):
     root = zarr.group(store=tmp_store)
     array = root.ones(name="ones", shape=(10, 10), chunks=(5, 5), dtype="float32")
@@ -31,7 +44,7 @@ async def test_pickle(tmp_store):
     store_loaded = pickle.loads(pickled)
     assert store_loaded == tmp_store
 
-    root_loaded = zarr.open_group(store_loaded)
+    root_loaded = zarr.open_group(store_loaded, mode="r")
     array_loaded = root_loaded["ones"]
 
     assert type(array_loaded) is zarr.Array