From 66397e89529b2a3ecaa84d4945bd00da12a573a7 Mon Sep 17 00:00:00 2001
From: Sam Levang <39069044+slevang@users.noreply.github.com>
Date: Mon, 27 Nov 2023 11:47:46 -0500
Subject: [PATCH] Keep attrs in `map_over_subtree` (#279)

* keep attrs in map_over_subtree

* more intelligible logic

---------

Co-authored-by: Tom Nicholas <tom@cworthy.org>
---
 datatree/mapping.py            | 15 +++++++++------
 datatree/tests/test_mapping.py | 11 +++++++++++
 docs/source/whats-new.rst      |  2 ++
 3 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/datatree/mapping.py b/datatree/mapping.py
index d6df73f1..c9631e1e 100644
--- 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
index 7ac5da5b..78486b06 100644
--- 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
index fe1ac57f..9d46c95b 100644
--- 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
 ~~~~~~~~~~~~~