def ToControlFlowGraph(g: nx.MultiDiGraph):
"""Create a new graph with only the statements and control flow edges."""
# CFGs cannot have parallel edges, so we use only a DiGraph rather than
# MultiDiGraph.
cfg = nx.DiGraph()
for node, data in nx_utils.StatementNodeIterator(g):
cfg.add_node(node, **data)
for src, dst, data in nx_utils.ControlFlowEdgeIterator(g):
cfg.add_edge(src, dst, **data)
def test_ControlFlowEdgeIterator(graph):
assert len(list(nx_utils.ControlFlowEdgeIterator(graph))) == 3