Skip to content

Commit 2f3b6fa

Browse files
committed
feat(net): find_op_XXX() utilitiy methods
1 parent 2fbfe43 commit 2f3b6fa

File tree

1 file changed

+20
-0
lines changed

1 file changed

+20
-0
lines changed

graphtik/network.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -940,6 +940,26 @@ def __repr__(self):
940940
)
941941
return f"Network(x{len(self.graph.nodes)} nodes, x{len(ops)} ops: {''.join(steps)})"
942942

943+
def find_ops(self, predicate) -> List[Operation]:
944+
"""
945+
Scan operation nodes and fetch those satisfying `predicate`.
946+
947+
:param predicate:
948+
the :term:`node predicate` is a 2-argument callable(op, node-data)
949+
that should return true for nodes to include.
950+
"""
951+
return [
952+
n
953+
for n, d in self.graph.nodes.data(True)
954+
if isinstance(n, Operation) and predicate(n, d)
955+
]
956+
957+
def find_op_by_name(self, name) -> Union[Operation, None]:
958+
"""Fetch the 1st operation named with the given `name`."""
959+
for n in yield_ops(self.graph.nodes):
960+
if n.name == name:
961+
return n
962+
943963
def prepare_plot_args(self, plot_args: PlotArgs) -> PlotArgs:
944964
plot_args = plot_args.clone_or_merge_graph(self.graph)
945965
plot_args = plot_args.with_defaults(

0 commit comments

Comments
 (0)