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

Add match_names function to filter specific names #343

Closed
wants to merge 1 commit into from
Closed
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
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
Loading