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

Commit

Permalink
Map over only data nodes, ignoring attrs (#263)
Browse files Browse the repository at this point in the history
* add test from issue

* test as a property of map_over_subtree directly

* change behaviour

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* correct test

* whatsnew

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
TomNicholas and pre-commit-ci[bot] authored Oct 24, 2023
1 parent 7ba0588 commit 5330a3c
Show file tree
Hide file tree
Showing 4 changed files with 25 additions and 3 deletions.
6 changes: 3 additions & 3 deletions datatree/mapping.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,8 +109,8 @@ def map_over_subtree(func: Callable) -> Callable:
Applies a function to every dataset in one or more subtrees, returning new trees which store the results.
The function will be applied to any non-empty dataset stored in any of the nodes in the trees. The returned trees
will have the same structure as the supplied trees.
The function will be applied to any data-containing dataset stored in any of the nodes in the trees. The returned
trees will have the same structure as the supplied trees.
`func` needs to return one Datasets, DataArrays, or None in order to be able to rebuild the subtrees after
mapping, as each result will be assigned to its respective node of a new tree via `DataTree.__setitem__`. Any
Expand Down Expand Up @@ -206,7 +206,7 @@ def _map_over_subtree(*args, **kwargs) -> DataTree | Tuple[DataTree, ...]:
# Now we can call func on the data in this particular set of corresponding nodes
results = (
func(*node_args_as_datasets, **node_kwargs_as_datasets)
if not node_of_first_tree.is_empty
if node_of_first_tree.has_data
else None
)

Expand Down
7 changes: 7 additions & 0 deletions datatree/tests/test_datatree.py
Original file line number Diff line number Diff line change
Expand Up @@ -603,6 +603,13 @@ def test_ipython_key_completions(self, create_test_datatree):
var_keys = list(dt.variables.keys())
assert all(var_key in key_completions for var_key in var_keys)

def test_operation_with_attrs_but_no_data(self):
# tests bug from xarray-datatree GH262
xs = xr.Dataset({"testvar": xr.DataArray(np.ones((2, 3)))})
dt = DataTree.from_dict({"node1": xs, "node2": xs})
dt.attrs["test_key"] = 1 # sel works fine without this line
dt.sel(dim_0=0)


class TestRestructuring:
def test_drop_nodes(self):
Expand Down
12 changes: 12 additions & 0 deletions datatree/tests/test_mapping.py
Original file line number Diff line number Diff line change
Expand Up @@ -252,6 +252,18 @@ def times_ten(ds):
result_tree = times_ten(subtree)
assert_equal(result_tree, expected, from_root=False)

def test_skip_empty_nodes_with_attrs(self, create_test_datatree):
# inspired by xarray-datatree GH262
dt = create_test_datatree()
dt["set1/set2"].attrs["foo"] = "bar"

def check_for_data(ds):
# fails if run on a node that has no data
assert len(ds.variables) != 0
return ds

dt.map_over_subtree(check_for_data)


class TestMutableOperations:
def test_construct_using_type(self):
Expand Down
3 changes: 3 additions & 0 deletions docs/source/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,9 @@ New Features
Breaking changes
~~~~~~~~~~~~~~~~

- Nodes containing only attributes but no data are now ignored by :py:func:`map_over_subtree` (:issue:`262`, :pull:`263`)
By `Tom Nicholas <https://github.com/TomNicholas>`_.

Deprecations
~~~~~~~~~~~~

Expand Down

0 comments on commit 5330a3c

Please sign in to comment.