Update datatree notebook (#188)

* add watermark

* fix return type in load_arviz_data

* update tutorial

Author: symeneses

docs/source/tutorial/WorkingWithDataTree.ipynb

@@ -66,7 +66,8 @@ doc = [
     "sphinx>=5",
     "sphinx-design",
     "jupyter-sphinx",
-    "h5netcdf[h5py]"
+    "h5netcdf[h5py]",
+    "watermark"
 ]
 netcdf4 = [
   "netcdf4"

pyproject.toml

@@ -112,7 +112,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.py

@@ -10,7 +10,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
 
@@ -33,5 +33,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]: ...