DataTree: sel & isel add error context (#10154)

* DataTree: sel & isel add error context

* add test

* changelog

---------

Co-authored-by: Deepak Cherian <[email protected]>

Author: mathause

doc/whats-new.rst

@@ -43,6 +43,8 @@ Bug fixes
   This fixes the issue where using the ``zarr_version`` parameter would raise a deprecation warning telling the user to use
   a non-existent ``zarr_format`` parameter instead. (:issue:`10163`, :pull:`10164`)
   By `Karl Krauth <https://github.com/Karl-Krauth>`_.
+- :py:meth:`DataTree.sel` and :py:meth:`DataTree.isel` display the path of the first failed node again  (:pull:`10154`).
+  By `Mathias Hauser <https://github.com/mathause>`_.
 - Fix grouped and resampled ``first``, ``last`` with datetimes (:issue:`10169`, :pull:`10173`)
   By `Deepak Cherian <https://github.com/dcherian>`_.
 

xarray/core/datatree.py

@@ -32,6 +32,7 @@
 from xarray.core.dataset import Dataset
 from xarray.core.dataset_variables import DataVariables
 from xarray.core.datatree_mapping import (
+    _handle_errors_with_path_context,
     map_over_datasets,
 )
 from xarray.core.formatting import (
@@ -1845,7 +1846,8 @@ def _selective_indexing(
         result = {}
         for path, node in self.subtree_with_keys:
             node_indexers = {k: v for k, v in indexers.items() if k in node.dims}
-            node_result = func(node.dataset, node_indexers)
+            func_with_error_context = _handle_errors_with_path_context(path)(func)
+            node_result = func_with_error_context(node.dataset, node_indexers)
             # Indexing datasets corresponding to each node results in redundant
             # coordinates when indexes from a parent node are inherited.
             # Ideally, we would avoid creating such coordinates in the first

xarray/tests/test_datatree.py

@@ -1955,6 +1955,26 @@ def test_sel(self) -> None:
         )
         assert_identical(actual, expected)
 
+    def test_sel_isel_error_has_node_info(self) -> None:
+        tree = DataTree.from_dict(
+            {
+                "/first": xr.Dataset({"a": ("x", [1, 2, 3])}, coords={"x": [1, 2, 3]}),
+                "/second": xr.Dataset({"b": ("x", [4, 5])}, coords={"x": [2, 3]}),
+            }
+        )
+
+        with pytest.raises(
+            KeyError,
+            match="Raised whilst mapping function over node with path 'second'",
+        ):
+            tree.sel(x=1)
+
+        with pytest.raises(
+            IndexError,
+            match="Raised whilst mapping function over node with path 'first'",
+        ):
+            tree.isel(x=4)
+
 
 class TestAggregations:
     def test_reduce_method(self) -> None: