Skip to content
This repository has been archived by the owner on Oct 24, 2024. It is now read-only.

Commit

Permalink
Keep attrs in map_over_subtree (#279)
Browse files Browse the repository at this point in the history
* keep attrs in map_over_subtree

* more intelligible logic

---------

Co-authored-by: Tom Nicholas <tom@cworthy.org>
  • Loading branch information
slevang and TomNicholas authored Nov 27, 2023
1 parent b62b6bc commit 66397e8
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 6 deletions.
15 changes: 9 additions & 6 deletions datatree/mapping.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
11 changes: 11 additions & 0 deletions datatree/tests/test_mapping.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)
Expand Down
2 changes: 2 additions & 0 deletions docs/source/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
~~~~~~~~~~~~~
Expand Down

0 comments on commit 66397e8

Please sign in to comment.