also apply fix in validate_partition_objects

Deltares · Aug 20, 2024 · b66b6b1 · b66b6b1
1 parent daed521
commit b66b6b1
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/xugrid/ugrid/conventions.py b/xugrid/ugrid/conventions.py
@@ -181,10 +181,10 @@ def default_topology_attrs(name: str, topology_dimension: int):
 
 
 def _get_topology(ds: xr.Dataset) -> List[str]:
-    # looping over ds.variables.items() is way faster for datasets with many 
-    # variables than accessing each var with ds[var]
+    # accessing ds.variables[var] is way faster for datasets with many 
+    # variables than accessing them with ds[var]
     # https://github.com/Deltares/xugrid/issues/285
-    return [varname for varname, var in ds.variables.items() if var.attrs.get("cf_role") == "mesh_topology"]
+    return [k for k in ds.data_vars if ds.variables[k].attrs.get("cf_role") == "mesh_topology"]
 
 
 def _infer_xy_coords(

diff --git a/xugrid/ugrid/partitioning.py b/xugrid/ugrid/partitioning.py
@@ -202,7 +202,7 @@ def validate_partition_objects(
         unique_vars = set(chain(*allvars))
         # Check dimensions
         dims_per_var = [
-            {ds[var].dims for ds in data_objects if var in ds.data_vars}
+            {ds.variables[var].dims for ds in data_objects if var in ds.data_vars}
             for var in unique_vars
         ]
         for var, vardims in zip(unique_vars, dims_per_var):