fix data catalog deprecation warnings

Deltares · Feb 13, 2025 · 402495a · 402495a
1 parent 20f5e87
commit 402495a
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 9 deletions.
diff --git a/hydromt_wflow/wflow.py b/hydromt_wflow/wflow.py
@@ -218,7 +218,7 @@ def setup_basemaps(
         xy = None
         if kind in ["basin", "subbasin", "outlet"]:
             if basin_index_fn is not None:
-                bas_index = self.data_catalog[basin_index_fn]
+                bas_index = self.data_catalog.get_source(basin_index_fn)
             else:
                 bas_index = None
             geom, xy = hydromt.workflows.get_basin_geometry(
@@ -1526,8 +1526,8 @@ def setup_gauges(
                 handle_nodata=NoDataStrategy.IGNORE,
                 **kwargs,
             )
-        elif gauges_fn in self.data_catalog:
-            if self.data_catalog[gauges_fn].data_type == "GeoDataFrame":
+        elif self.data_catalog.contains_source(gauges_fn):
+            if self.data_catalog.get_source(gauges_fn).data_type == "GeoDataFrame":
                 gdf_gauges = self.data_catalog.get_geodataframe(
                     gauges_fn,
                     geom=self.basins,
@@ -1630,7 +1630,7 @@ def setup_gauges(
             return
 
         # Add to grid
-        mapname = f'{str(self._MAPS["gauges"])}_{basename}'
+        mapname = f"{str(self._MAPS['gauges'])}_{basename}"
         self.set_grid(da, name=mapname)
 
         # geoms
@@ -1815,7 +1815,9 @@ def setup_lakes(
                     i = fns_ids.index(id)
                     rating_fn = rating_curve_fns[i]
                     # Read data
-                    if isfile(rating_fn) or rating_fn in self.data_catalog:
+                    if isfile(rating_fn) or self.data_catalog.contains_source(
+                        rating_fn
+                    ):
                         self.logger.info(
                             f"Preparing lake rating curve data from {rating_fn}"
                         )
@@ -4363,7 +4365,7 @@ def set_grid(
                 )
             tname = "time"
             time_axes = {
-                k: v for k, v in dict(self.grid.dims).items() if k.startswith("time")
+                k: v for k, v in self.grid.sizes.items() if k.startswith("time")
             }
             if data["time"].size not in time_axes.values():
                 tname = f"time_{data['time'].size}" if "time" in time_axes else tname
@@ -4860,7 +4862,7 @@ def read_intbl(self, **kwargs):
                 name = basename(fn).split(".")[0]
                 tbl = pd.read_csv(fn, delim_whitespace=True, header=None)
                 tbl.columns = [
-                    f"expr{i+1}" if i + 1 < len(tbl.columns) else "value"
+                    f"expr{i + 1}" if i + 1 < len(tbl.columns) else "value"
                     for i in range(len(tbl.columns))
                 ]  # rename columns
                 self.set_intbl(tbl, name=name)

diff --git a/tests/test_model_methods.py b/tests/test_model_methods.py
@@ -277,8 +277,10 @@ def test_setup_reservoirs(source, tmpdir, example_wflow_model):
                 )
             )
             == number_of_reservoirs
-        ), f"Number of non-null values in {i} not equal to \
+        ), (
+            f"Number of non-null values in {i} not equal to \
 number of reservoirs in model area"
+        )
 
 
 def test_setup_ksathorfrac(tmpdir, example_wflow_model):
@@ -530,7 +532,7 @@ def test_setup_outlets(example_wflow_model):
 def test_setup_gauges(example_wflow_model):
     # 1. Test with grdc data
     # uparea rename not in the latest artifact_data version
-    example_wflow_model.data_catalog["grdc"].rename = {"area": "uparea"}
+    example_wflow_model.data_catalog.get_source("grdc").rename = {"area": "uparea"}
     example_wflow_model.setup_gauges(
         gauges_fn="grdc",
         basename="grdc_uparea",