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

Commit

Permalink
Add match_names function to filter specific names
Browse files Browse the repository at this point in the history
  • Loading branch information
Armavica committed Sep 7, 2024
1 parent d97b636 commit 12b280d
Show file tree
Hide file tree
Showing 4 changed files with 70 additions and 0 deletions.
46 changes: 46 additions & 0 deletions datatree/datatree.py
Original file line number Diff line number Diff line change
Expand Up @@ -1307,6 +1307,52 @@ 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"])
DataTree('None', parent=None)
├── DataTree('a')
│ └── DataTree('A')
│ └── DataTree('C')
└── DataTree('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 datatree/tests/test_datatree.py
Original file line number Diff line number Diff line change
Expand Up @@ -707,6 +707,27 @@ def test_match(self):
)
dtt.assert_identical(result, expected)

def test_match_names(self):
# TODO is this example going to cause problems with case sensitivity?
dt = 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,
}
)
dtt.assert_identical(result, expected)

def test_filter(self):
simpsons = DataTree.from_dict(
d={
Expand Down
1 change: 1 addition & 0 deletions docs/source/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,7 @@ For manipulating, traversing, navigating, or mapping over the tree structure.
map_over_subtree
DataTree.pipe
DataTree.match
DataTree.match_names
DataTree.filter

Pathlib-like Interface
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 @@ -23,6 +23,8 @@ v0.0.14 (unreleased)
New Features
~~~~~~~~~~~~

- Added `DataTree.match_names` method to filter a list of specific node names.

Breaking changes
~~~~~~~~~~~~~~~~

Expand Down

0 comments on commit 12b280d

Please sign in to comment.