fix return type in load_arviz_data

Author: symeneses

src/arviz_base/datasets.py

@@ -111,7 +111,7 @@ def load_arviz_data(dataset=None, data_home=None, **kwargs):
 
     Returns
     -------
-    xarray.Dataset
+    DataTree
     """
     if dataset in LOCAL_DATASETS:
         resource = LOCAL_DATASETS[dataset]

src/arviz_base/datasets.pyi

@@ -9,7 +9,7 @@ from collections import namedtuple
 from urllib.request import urlretrieve
 
 from _typeshed import Incomplete
-from xarray import Dataset, open_datatree
+from xarray import DataTree, open_datatree
 
 from arviz_base.rcparams import rcParams
 
@@ -32,5 +32,5 @@ def clear_data_home(data_home: str | None = ...) -> None: ...
 def _sha256(path) -> None: ...
 def load_arviz_data(
     dataset: str | None = ..., data_home: str | None = ..., **kwargs: dict
-) -> Dataset: ...
+) -> DataTree: ...
 def list_datasets() -> list[str]: ...