diff --git a/icechunk-python/python/icechunk/__init__.py b/icechunk-python/python/icechunk/__init__.py index e737a0ab..af56201f 100644 --- a/icechunk-python/python/icechunk/__init__.py +++ b/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) diff --git a/icechunk-python/tests/test_pickle.py b/icechunk-python/tests/test_pickle.py index e5b06f7c..774576fb 100644 --- a/icechunk-python/tests/test_pickle.py +++ b/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