ENH: enable H5NetCDFStore to work with already open h5netcdf.File and h5netcdf.Group objects, add test

Author: kmuehlbauer

xarray/backends/api.py

@@ -503,7 +503,7 @@ def maybe_decode_store(store, lock=False):
         elif engine == "pydap":
             store = backends.PydapDataStore.open(filename_or_obj, **backend_kwargs)
         elif engine == "h5netcdf":
-            store = backends.H5NetCDFStore(
+            store = backends.H5NetCDFStore.open(
                 filename_or_obj, group=group, lock=lock, **backend_kwargs
             )
         elif engine == "pynio":
@@ -968,7 +968,7 @@ def open_mfdataset(
 WRITEABLE_STORES: Dict[str, Callable] = {
     "netcdf4": backends.NetCDF4DataStore.open,
     "scipy": backends.ScipyDataStore,
-    "h5netcdf": backends.H5NetCDFStore,
+    "h5netcdf": backends.H5NetCDFStore.open,
 }
 
 

xarray/backends/h5netcdf_.py

@@ -4,9 +4,9 @@
 
 from .. import Variable
 from ..core import indexing
-from ..core.utils import FrozenDict
+from ..core.utils import FrozenDict, is_remote_uri
 from .common import WritableCFDataStore
-from .file_manager import CachingFileManager
+from .file_manager import CachingFileManager, DummyFileManager
 from .locks import HDF5_LOCK, combine_locks, ensure_lock, get_write_lock
 from .netCDF4_ import (
     BaseNetCDF4Array,
@@ -69,8 +69,42 @@ class H5NetCDFStore(WritableCFDataStore):
     """Store for reading and writing data via h5netcdf
     """
 
+    __slots__ = (
+        "autoclose",
+        "format",
+        "is_remote",
+        "lock",
+        "_filename",
+        "_group",
+        "_manager",
+        "_mode",
+    )
+
     def __init__(
-        self,
+            self, manager, group=None, mode=None, lock=HDF5_LOCK, autoclose=False
+    ):
+
+        import h5netcdf
+
+        if isinstance(manager, (h5netcdf.File, h5netcdf.Group)):
+            if group is None:
+                root, group = manager._root, manager.name
+            else:
+                root = manager
+            manager = DummyFileManager(root)
+
+        self._manager = manager
+        self._group = group
+        self._mode = mode
+        self.format = None
+        self._filename = self.ds._root.filename
+        self.is_remote = is_remote_uri(self._filename)
+        self.lock = ensure_lock(lock)
+        self.autoclose = autoclose
+
+    @classmethod
+    def open(
+        cls,
         filename,
         mode="r",
         format=None,
@@ -86,22 +120,16 @@ def __init__(
 
         kwargs = {"invalid_netcdf": invalid_netcdf}
 
-        self._manager = CachingFileManager(
-            h5netcdf.File, filename, mode=mode, kwargs=kwargs
-        )
-
         if lock is None:
             if mode == "r":
                 lock = HDF5_LOCK
             else:
                 lock = combine_locks([HDF5_LOCK, get_write_lock(filename)])
 
-        self._group = group
-        self.format = format
-        self._filename = filename
-        self._mode = mode
-        self.lock = ensure_lock(lock)
-        self.autoclose = autoclose
+        manager = CachingFileManager(
+            h5netcdf.File, filename, mode=mode, kwargs=kwargs
+        )
+        return cls(manager, group=group, mode=mode, lock=lock, autoclose=autoclose)
 
     def _acquire(self, needs_lock=True):
         with self._manager.acquire_context(needs_lock) as root:

xarray/tests/test_backends.py

@@ -2182,7 +2182,7 @@ class TestH5NetCDFData(NetCDF4Base):
     @contextlib.contextmanager
     def create_store(self):
         with create_tmp_file() as tmp_file:
-            yield backends.H5NetCDFStore(tmp_file, "w")
+            yield backends.H5NetCDFStore.open(tmp_file, "w")
 
     @pytest.mark.filterwarnings("ignore:complex dtypes are supported by h5py")
     @pytest.mark.parametrize(
@@ -2345,6 +2345,26 @@ def test_dump_encodings_h5py(self):
             assert actual.x.encoding["compression"] == "lzf"
             assert actual.x.encoding["compression_opts"] is None
 
+    def test_already_open_dataset_group(self):
+        import h5netcdf
+        with create_tmp_file() as tmp_file:
+            with nc4.Dataset(tmp_file, mode="w") as nc:
+                group = nc.createGroup("g")
+                v = group.createVariable("x", "int")
+                v[...] = 42
+
+            h5 = h5netcdf.File(tmp_file, mode="r")
+            store = backends.H5NetCDFStore(h5["g"])
+            with open_dataset(store) as ds:
+                expected = Dataset({"x": ((), 42)})
+                assert_identical(expected, ds)
+
+            h5 = h5netcdf.File(tmp_file, mode="r")
+            store = backends.H5NetCDFStore(h5, group="g")
+            with open_dataset(store) as ds:
+                expected = Dataset({"x": ((), 42)})
+                assert_identical(expected, ds)
+
 
 @requires_h5netcdf
 class TestH5NetCDFFileObject(TestH5NetCDFData):