Skip to content

Commit

Permalink
Add DataTree.match_names to match node names
Browse files Browse the repository at this point in the history
  • Loading branch information
Armavica committed Sep 7, 2024
1 parent 25debff commit 94e93fe
Show file tree
Hide file tree
Showing 2 changed files with 66 additions and 0 deletions.
45 changes: 45 additions & 0 deletions xarray/core/datatree.py
Original file line number Diff line number Diff line change
Expand Up @@ -1349,6 +1349,51 @@ def match(self, pattern: str) -> DataTree:
}
return DataTree.from_dict(matching_nodes, name=self.root.name)

def match_names(self, names: Iterable[str]) -> DataTree:
"""
Filter nodes by name.
Parameters
----------
names: Iterable[str]
The list of node names to retain.
Returns
-------
DataTree
See Also
--------
match
filter
pipe
map_over_subtree
Examples
--------
>>> dt = DataTree.from_dict(
... {
... "/a/A": None,
... "/a/B": None,
... "/a/C": None,
... "/C/D": None,
... "/E/F": None,
... }
... )
>>> dt.match_names(["A", "C"])
<xarray.DataTree>
Group: /
├── Group: /C
└── Group: /a
├── Group: /a/A
└── Group: /a/C
"""
names = set(names)
matching_nodes = {
node.path: node.ds for node in self.subtree if node.name in names
}
return DataTree.from_dict(matching_nodes, name=self.root.name)

def map_over_subtree(
self,
func: Callable,
Expand Down
21 changes: 21 additions & 0 deletions xarray/tests/test_datatree.py
Original file line number Diff line number Diff line change
Expand Up @@ -1025,6 +1025,27 @@ def test_match(self):
)
assert_identical(result, expected)

def test_match_names(self):
# TODO is this example going to cause problems with case sensitivity?
dt: DataTree = DataTree.from_dict(
{
"/a/A": None,
"/a/B": None,
"/a/C": None,
"/C/D": None,
"/E/F": None,
}
)
result = dt.match_names(["A", "C"])
expected = DataTree.from_dict(
{
"/a/A": None,
"/a/C": None,
"/C": None,
}
)
assert_identical(result, expected)

def test_filter(self):
simpsons: DataTree = DataTree.from_dict(
d={
Expand Down

0 comments on commit 94e93fe

Please sign in to comment.