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

Keep attrs in map_over_subtree #279

Merged
merged 3 commits into from
Nov 27, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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