Fix type checking consistency

Author: jenshnielsen

src/qcodes/dataset/data_set.py

@@ -961,7 +961,7 @@ def to_xarray_dataarray_dict(
         start: int | None = None,
         end: int | None = None,
         use_multi_index: Literal["auto", "always", "never"] = "auto",
-    ) -> dict[str, xr.DataArray]:
+    ) -> dict[str, xr.Dataset]:
         """
         Returns the values stored in the :class:`.DataSet` for the specified parameters
         and their dependencies as a dict of :py:class:`xr.DataArray` s

src/qcodes/dataset/data_set_cache.py

@@ -187,7 +187,7 @@ def to_pandas_dataframe(self) -> pd.DataFrame:
 
     def to_xarray_dataarray_dict(
         self, *, use_multi_index: Literal["auto", "always", "never"] = "auto"
-    ) -> dict[str, xr.DataArray]:
+    ) -> dict[str, xr.Dataset]:
         """
         Returns the values stored in the :class:`.dataset.data_set.DataSet` as a dict of
         :py:class:`xr.DataArray` s

src/qcodes/dataset/data_set_protocol.py

@@ -231,7 +231,7 @@ def to_xarray_dataarray_dict(
         start: int | None = None,
         end: int | None = None,
         use_multi_index: Literal["auto", "always", "never"] = "auto",
-    ) -> dict[str, xr.DataArray]: ...
+    ) -> dict[str, xr.Dataset]: ...
 
     def to_xarray_dataset(
         self,