Keep attrs in map_over_subtree xarray-contrib/datatree#279

* keep attrs in map_over_subtree

* more intelligible logic

---------

Co-authored-by: Tom Nicholas <[email protected]>

Author: slevang

datatree/mapping.py

@@ -206,14 +206,17 @@ def _map_over_subtree(*args, **kwargs) -> DataTree | Tuple[DataTree, ...]:
                 node_of_first_tree.path
             )(func)
 
-            # Now we can call func on the data in this particular set of corresponding nodes
-            results = (
-                func_with_error_context(
+            if node_of_first_tree.has_data:
+                # call func on the data in this particular set of corresponding nodes
+                results = func_with_error_context(
                     *node_args_as_datasetviews, **node_kwargs_as_datasetviews
                 )
-                if node_of_first_tree.has_data
-                else None
-            )
+            elif node_of_first_tree.has_attrs:
+                # propagate attrs
+                results = node_of_first_tree.ds
+            else:
+                # nothing to propagate so use fastpath to create empty node in new tree
+                results = None
 
             # TODO implement mapping over multiple trees in-place using if conditions from here on?
             out_data_objects[node_of_first_tree.path] = results

datatree/tests/test_mapping.py

@@ -264,6 +264,17 @@ def check_for_data(ds):
 
         dt.map_over_subtree(check_for_data)
 
+    def test_keep_attrs_on_empty_nodes(self, create_test_datatree):
+        # GH278
+        dt = create_test_datatree()
+        dt["set1/set2"].attrs["foo"] = "bar"
+
+        def empty_func(ds):
+            return ds
+
+        result = dt.map_over_subtree(empty_func)
+        assert result["set1/set2"].attrs == dt["set1/set2"].attrs
+
     @pytest.mark.xfail(
         reason="probably some bug in pytests handling of exception notes"
     )

docs/source/whats-new.rst

@@ -35,6 +35,8 @@ Deprecations
 
 Bug fixes
 ~~~~~~~~~
+- Keep attributes on nodes containing no data in :py:func:`map_over_subtree`. (:issue:`278`, :pull:`279`)
+  By `Sam Levang <https://github.com/slevang>`_.
 
 Documentation
 ~~~~~~~~~~~~~