Keep attrs in map_over_subtree (#279)

* keep attrs in map_over_subtree * more intelligible logic --------- Co-authored-by: Tom Nicholas <[email protected]>
xarray-contrib · Nov 27, 2023 · 66397e8 · 66397e8
1 parent b62b6bc
commit 66397e8
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 6 deletions.
diff --git a/datatree/mapping.py b/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

diff --git a/datatree/tests/test_mapping.py b/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"
     )

diff --git a/docs/source/whats-new.rst b/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
 ~~~~~~~~~~~~~