[bug] ds.variables and ds.data_vars are different things

Author: nikeethr

src/scores/typing.py

@@ -112,7 +112,7 @@ def make_dataarray_ref(self) -> xr.DataArray:
         if not self.is_dataarray():
             raise TypeError("Cannot revert name for xr.Dataset, only xr.DataArray")
 
-        var_names: list[Hashable] = list(self.ds.variables.keys())
+        var_names: list[Hashable] = list(self.ds.data_vars.keys())
 
         # safety: we should only have 1 variable if this is actually a data array
         assert len(var_names) == 1

tests/continuous/test_nse.py

@@ -846,7 +846,7 @@ def test_nse_with_datasets(self):
         assert set(res["temp"].dims) == set(["t", "h"])
         assert res["temp"].shape == (2, 4)
         # tapioca is ignored
-        assert "tapioca" not in res.variables.keys()
+        assert "tapioca" not in res.data_vars.keys()
 
 
 class TestNseDask(NseSetup):

tests/test_typing.py

@@ -32,7 +32,7 @@ def test_lift_dataarray_withname():
     da = xr.DataArray([1], dims="t", name=da_test_name)
     lds = scores.typing.LiftedDataset(da)
     # check lifted dataset only has one array
-    keys = list(lds.ds.variables.keys())
+    keys = list(lds.ds.data_vars.keys())
     assert len(keys) == 1
     # check that the name is preserved and isn't a dummy
     assert keys[0] == da_test_name
@@ -59,7 +59,7 @@ def test_lift_dataarray_withnameequalsdummy():
     da = xr.DataArray([1], dims="t", name=da_test_name)
     lds = scores.typing.LiftedDataset(da)
     # check lifted dataset only has one array
-    keys = list(lds.ds.variables.keys())
+    keys = list(lds.ds.data_vars.keys())
     assert len(keys) == 1
     # check that the name is the same as the dummy, BUT it isn't actually flagged as a dummy
     assert keys[0] == scores.typing.LiftedDataset._DUMMY_DATAARRAY_VARNAME  # pylint: disable-msg=protected-access
@@ -82,7 +82,7 @@ def test_lift_dataarray_noname():
     da = xr.DataArray([1], dims="t")
     lds = scores.typing.LiftedDataset(da)
     # check lifted dataset only has one array
-    keys = list(lds.ds.variables.keys())
+    keys = list(lds.ds.data_vars.keys())
     assert len(keys) == 1
     # check that the name is the same as the dummy, AND is also flagged as dummy
     assert keys[0] == scores.typing.LiftedDataset._DUMMY_DATAARRAY_VARNAME  # pylint: disable-msg=protected-access