Add path to error message in map_over_subtree

datatree/mapping.py

@@ -202,10 +202,15 @@ def _map_over_subtree(*args, **kwargs) -> DataTree | Tuple[DataTree, ...]:
                     ],
                 )
             )
+            func_with_error_context = _handle_errors_with_path_context(
+                node_of_first_tree.path
+            )(func)
 
             # Now we can call func on the data in this particular set of corresponding nodes
             results = (
-                func(*node_args_as_datasets, **node_kwargs_as_datasets)
+                func_with_error_context(
+                    *node_args_as_datasets, **node_kwargs_as_datasets
+                )
                 if not node_of_first_tree.is_empty
                 else None
             )
@@ -251,6 +256,23 @@ def _map_over_subtree(*args, **kwargs) -> DataTree | Tuple[DataTree, ...]:
     return _map_over_subtree
 
 
+def _handle_errors_with_path_context(path):
+    """Wraps given function so that if it fails it also raises path to node on which it failed."""
+
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except Exception as e:
+                # Add the context information to the error message
+                e.add_note(f"Raised whilst mapping function over node with path {path}")
+                raise
+
+        return wrapper
+
+    return decorator
+
+
 def _check_single_set_return_values(path_to_node, obj):
     """Check types returned from single evaluation of func, and return number of return values received from func."""
     if isinstance(obj, (Dataset, DataArray)):

datatree/tests/test_mapping.py

@@ -252,6 +252,20 @@ def times_ten(ds):
         result_tree = times_ten(subtree)
         assert_equal(result_tree, expected, from_root=False)
 
+    def test_error_contains_path_of_offending_node(self, create_test_datatree):
+        dt = create_test_datatree()
+        dt["set1"]["bad_var"] = 0
+        print(dt)
+
+        def fail_on_specific_node(ds):
+            if "bad_var" in ds:
+                raise ValueError("Failed because 'bar_var' present in dataset")
+
+        with pytest.raises(
+            ValueError, match="Raised whilst mapping function over node /set1"
+        ):
+            dt.map_over_subtree(fail_on_specific_node)
+
 
 class TestMutableOperations:
     def test_construct_using_type(self):