Support nested allow_pickling contexts (#790)

Co-authored-by: Matthew Iannucci <matthew@earthmover.io>

icechunk-python/python/icechunk/session.py

@@ -126,11 +126,14 @@ def allow_pickling(self) -> Generator[None, None, None]:
         """
         Context manager to allow unpickling this store if writable.
         """
+        # While this property can only be changed by this context manager,
+        # it can be nested (sometimes unintentionally since `to_icechunk` does it)
+        current = self._allow_pickling
         try:
             self._allow_pickling = True
             yield
         finally:
-            self._allow_pickling = False
+            self._allow_pickling = current
 
     @property
     def read_only(self) -> bool:

icechunk-python/tests/test_dask.py

@@ -3,8 +3,13 @@
 pytest.importorskip("dask")
 pytest.importorskip("distributed")
 
+import tempfile
+
 import dask
 import distributed
+import xarray as xr
+from icechunk import Repository, local_filesystem_storage
+from icechunk.xarray import to_icechunk
 from tests.test_xarray import create_test_data, roundtrip
 from xarray.testing import assert_identical
 
@@ -21,3 +26,22 @@ def test_threaded() -> None:
         ds = create_test_data().chunk(dim1=3, dim2=4)
         with roundtrip(ds) as actual:
             assert_identical(actual, ds)
+
+
+def test_xarray_to_icechunk_nested_pickling() -> None:
+    with dask.config.set(scheduler="processes"):
+        ds = create_test_data(dim_sizes=(2, 3, 4)).chunk(-1)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            repo = Repository.create(local_filesystem_storage(tmpdir))
+            session = repo.writable_session("main")
+
+            with session.allow_pickling():
+                to_icechunk(ds, session=session, mode="w")
+                with xr.open_zarr(session.store, consolidated=False) as actual:
+                    assert_identical(actual, ds)
+
+        newds = ds + 1
+        to_icechunk(newds, session=session, mode="w")
+        with session.allow_pickling():
+            with xr.open_zarr(session.store, consolidated=False) as actual:
+                assert_identical(actual, newds)