Do not call __exit__ on Zarr store when opening

datatree/io.py

@@ -84,24 +84,24 @@ def _open_datatree_netcdf(filename: str, **kwargs) -> DataTree:
 def _open_datatree_zarr(store, **kwargs) -> DataTree:
     import zarr  # type: ignore
 
-    with zarr.open_group(store, mode="r") as zds:
-        ds = open_dataset(store, engine="zarr", **kwargs)
-        tree_root = DataTree.from_dict({"/": ds})
-        for path in _iter_zarr_groups(zds):
-            try:
-                subgroup_ds = open_dataset(store, engine="zarr", group=path, **kwargs)
-            except zarr.errors.PathNotFoundError:
-                subgroup_ds = Dataset()
-
-            # TODO refactor to use __setitem__ once creation of new nodes by assigning Dataset works again
-            node_name = NodePath(path).name
-            new_node: DataTree = DataTree(name=node_name, data=subgroup_ds)
-            tree_root._set_item(
-                path,
-                new_node,
-                allow_overwrite=False,
-                new_nodes_along_path=True,
-            )
+    zds = zarr.open_group(store, mode="r")
+    ds = open_dataset(store, engine="zarr", **kwargs)
+    tree_root = DataTree.from_dict({"/": ds})
+    for path in _iter_zarr_groups(zds):
+        try:
+            subgroup_ds = open_dataset(store, engine="zarr", group=path, **kwargs)
+        except zarr.errors.PathNotFoundError:
+            subgroup_ds = Dataset()
+
+        # TODO refactor to use __setitem__ once creation of new nodes by assigning Dataset works again
+        node_name = NodePath(path).name
+        new_node: DataTree = DataTree(name=node_name, data=subgroup_ds)
+        tree_root._set_item(
+            path,
+            new_node,
+            allow_overwrite=False,
+            new_nodes_along_path=True,
+        )
     return tree_root
 
 

datatree/tests/test_io.py

@@ -40,6 +40,20 @@ def test_to_zarr(self, tmpdir):
         roundtrip_dt = open_datatree(filepath, engine="zarr")
         assert_equal(original_dt, roundtrip_dt)
 
+    @requires_zarr
+    def test_to_zarr_zip_store(self, tmpdir):
+        from zarr.storage import ZipStore
+
+        filepath = str(
+            tmpdir / "test.zarr.zip"
+        )  # casting to str avoids a pathlib bug in xarray
+        original_dt = create_test_datatree()
+        store = ZipStore(filepath)
+        original_dt.to_zarr(store)
+
+        roundtrip_dt = open_datatree(store, engine="zarr")
+        assert_equal(original_dt, roundtrip_dt)
+
     @requires_zarr
     def test_to_zarr_not_consolidated(self, tmpdir):
         filepath = tmpdir / "test.zarr"