From b66b6b1d7492a920933fcf242d8f2110e4daa3ea Mon Sep 17 00:00:00 2001
From: veenstrajelmer <60435591+veenstrajelmer@users.noreply.github.com>
Date: Tue, 20 Aug 2024 15:06:35 +0200
Subject: [PATCH] also apply fix in validate_partition_objects

---
 xugrid/ugrid/conventions.py  | 6 +++---
 xugrid/ugrid/partitioning.py | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/xugrid/ugrid/conventions.py b/xugrid/ugrid/conventions.py
index dc6eab0b8..acdf87214 100644
--- 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
index 0709bfc73..4d8dad04a 100644
--- 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):