Skip to content
This repository has been archived by the owner on Oct 24, 2024. It is now read-only.

Commit

Permalink
Do not call __exit__ on Zarr store when opening
Browse files Browse the repository at this point in the history
The `with` context when opening the zarr group with result in calling
__exit__ on the store when the function completes. This calls `.close()`
on ZipStore's, which results in errors:

```
ValueError: Attempt to use ZIP archive that was already closed
```
  • Loading branch information
thewtex committed May 16, 2022
1 parent badb7bc commit bebf63a
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 18 deletions.
36 changes: 18 additions & 18 deletions datatree/io.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
14 changes: 14 additions & 0 deletions datatree/tests/test_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down

0 comments on commit bebf63a

Please sign in to comment.