Coverage for src/loman/graph_utils.py: 100%

1"""Graph utility functions for computation graph operations.""" 

2 

3import functools 

4 

5from loman.util import apply_n 

6 

7 

8def contract_node_one(g, n): 

9 """Remove a node from graph and connect its predecessors to its successors.""" 

10 for p in g.predecessors(n): 

11 for s in g.successors(n): 

12 g.add_edge(p, s) 

13 g.remove_node(n) 

14 

15 

16def contract_node(g, ns): 

17 """Remove multiple nodes from graph and connect their predecessors to successors.""" 

18 apply_n(functools.partial(contract_node_one, g), ns)