added xu_open_dataset wrapper (#236)

* added xu_open_dataset wrapper. * load_dataset -> open_dataset * linting --------- Co-authored-by: roeldegoede <[email protected]>
Deltares · Jan 15, 2025 · e776a8e · e776a8e
1 parent 92a6080
commit e776a8e
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 5 deletions.
diff --git a/hydromt_sfincs/quadtree.py b/hydromt_sfincs/quadtree.py
@@ -11,6 +11,8 @@
 import xugrid as xu
 from pyproj import CRS, Transformer
 
+from hydromt_sfincs.utils import xu_open_dataset
+
 # optional dependency
 try:
     import datashader.transfer_functions as tf
@@ -83,8 +85,7 @@ def empty_mask(self):
     def read(self, file_name: Union[str, Path] = "sfincs.nc"):
         """Reads a quadtree netcdf file and stores it in the QuadtreeGrid object."""
 
-        self.data = xu.open_dataset(file_name)
-        self.data.close()  # TODO check if close works/is needed
+        self.data = xu_open_dataset(file_name)
 
         # TODO make similar to fortran conventions?
         # Rename to python conventions

diff --git a/hydromt_sfincs/sfincs.py b/hydromt_sfincs/sfincs.py
@@ -1350,7 +1350,7 @@ def setup_cn_infiltration_with_ks(
                 ib += 1
                 self.logger.debug(
                     f"\nblock {ib + 1}/{nrbn * nrbm} -- "
-                    f"col {bm0}:{bm1-1} | row {bn0}:{bn1-1}"
+                    f"col {bm0}:{bm1 - 1} | row {bn0}:{bn1 - 1}"
                 )
 
                 # calculate transform and shape of block at cell and subgrid level
@@ -3189,21 +3189,26 @@ def read_forcing(self, data_vars: List = None):
         for name in dvars_2d:
             fname, rename = self._FORCING_NET[name]
             fn = self.get_config(f"{fname}file", abs_path=True)
+            ds = None
             if fn is None or not isfile(fn):
                 if fn is not None:
                     self.logger.warning(f"{name}file not found at {fn}")
                 continue
             elif name in ["netbndbzsbzi", "netsrcdis"]:
                 ds = GeoDataset.from_netcdf(fn, crs=self.crs, chunks="auto")
             else:
-                ds = xr.load_dataset(fn, chunks="auto")
+                ds = xr.open_dataset(fn, chunks="auto")
+
             rename = {k: v for k, v in rename.items() if k in ds}
             if len(rename) > 0:
                 ds = ds.rename(rename).squeeze(drop=True)[list(rename.values())]
                 self.set_forcing(ds, split_dataset=True)
             else:
                 logger.warning(f"No forcing variables found in {fname}file")
 
+            if ds is not None:
+                ds.close()
+
     def write_forcing(self, data_vars: Union[List, str] = None, fmt: str = "%7.2f"):
         """Write forcing to ascii or netcdf (netampr) files.
         Filenames are based on the `config` attribute.
@@ -3425,7 +3430,7 @@ def read_results(
                 self.set_results(ds_face, split_dataset=True)
                 self.set_results(ds_edge, split_dataset=True)
             elif self.grid_type == "quadtree":
-                dsu = xu.open_dataset(
+                dsu = utils.xu_open_dataset(
                     fn_map,
                     chunks={"time": chunksize},
                 )

diff --git a/hydromt_sfincs/utils.py b/hydromt_sfincs/utils.py
@@ -1329,3 +1329,12 @@ def binary_search(vals, val):
         if vals[indx] == val:
             return indx
     return None
+
+
+def xu_open_dataset(*args, **kwargs):
+    """This function is a replacement of xu.open_dataset.
+
+    It exists because xu.open_dataset does not close the file after opening, which can lead to Permission Errors.
+    """
+    with xr.open_dataset(*args, **kwargs) as ds:
+        return xu.UgridDataset(ds)