Fix type checking consistency

Author: jenshnielsen

src/qcodes/dataset/data_set_cache.py

@@ -200,9 +200,10 @@ def to_xarray_dataarray_dict(
 
         """
         data = self.data()
-        return load_to_xarray_dataarray_dict(
+        data_dict = load_to_xarray_dataarray_dict(
             self._dataset, data, use_multi_index=use_multi_index
         )
+        return {key: val[key] for key, val in data_dict.items()}
 
     def to_xarray_dataset(
         self, *, use_multi_index: Literal["auto", "always", "never"] = "auto"

src/qcodes/dataset/exporters/export_to_xarray.py

@@ -61,7 +61,7 @@ def _calculate_index_shape(idx: pd.Index | pd.MultiIndex) -> dict[Hashable, int]
     return expanded_shape
 
 
-def _load_to_xarray_dataarray_dict_no_metadata(
+def _load_to_xarray_dataset_dict_no_metadata(
     dataset: DataSetProtocol,
     datadict: Mapping[str, Mapping[str, npt.NDArray]],
     *,
@@ -220,18 +220,21 @@ def load_to_xarray_dataarray_dict(
     datadict: Mapping[str, Mapping[str, npt.NDArray]],
     *,
     use_multi_index: Literal["auto", "always", "never"] = "auto",
-) -> dict[str, xr.Dataset]:
-    dataarrays = _load_to_xarray_dataarray_dict_no_metadata(
+) -> dict[str, xr.DataArray]:
+    xr_datasets = _load_to_xarray_dataset_dict_no_metadata(
         dataset, datadict, use_multi_index=use_multi_index
     )
+    data_arrays: dict[str, xr.DataArray] = {}
 
-    for dataname, dataarray in dataarrays.items():
-        _add_param_spec_to_xarray_coords(dataset, dataarray)
+    for dataname, xr_dataset in xr_datasets.items():
+        data_array = xr_dataset[dataname]
+        _add_param_spec_to_xarray_coords(dataset, data_array)
         paramspec_dict = _paramspec_dict_with_extras(dataset, str(dataname))
-        dataarray.attrs.update(paramspec_dict.items())
-        _add_metadata_to_xarray(dataset, dataarray)
+        data_array.attrs.update(paramspec_dict.items())
+        _add_metadata_to_xarray(dataset, data_array)
+        data_arrays[dataname] = data_array
 
-    return dataarrays
+    return data_arrays
 
 
 def _add_metadata_to_xarray(
@@ -270,7 +273,7 @@ def load_to_xarray_dataset(
 ) -> xr.Dataset:
     import xarray as xr
 
-    data_xrdarray_dict = _load_to_xarray_dataarray_dict_no_metadata(
+    data_xrdarray_dict = _load_to_xarray_dataset_dict_no_metadata(
         dataset, data, use_multi_index=use_multi_index
     )