def py_ast_to_graph(tree):
"""Convert an unsimplified AST into a graph, with a forward mapping.
Args:
tree: The (unsimplified) AST for the program.
Returns:
- Graph representing the tree.
- Dictionary that maps from AST node ids to graph node ids.
"""
return generic_ast_graphs.ast_to_graph(root=py_ast_to_generic(tree),
ast_spec=PY_AST_SPECS,
ignore_unknown_fields=True)
def test_compute_nth_child_edges(self):
mini_ast_spec = {
"root":
generic_ast_graphs.ASTNodeSpec(
fields={"children": generic_ast_graphs.FieldType.SEQUENCE},
sequence_item_types={"children": "child"},
has_parent=False),
"leaf":
generic_ast_graphs.ASTNodeSpec()
}
mini_ast_node = generic_ast_graphs.GenericASTNode(
"root", "root", {
"children": [
generic_ast_graphs.GenericASTNode("leaf0", "leaf", {}),
generic_ast_graphs.GenericASTNode("leaf1", "leaf", {}),
generic_ast_graphs.GenericASTNode("leaf2", "leaf", {}),
generic_ast_graphs.GenericASTNode("leaf3", "leaf", {}),
generic_ast_graphs.GenericASTNode("leaf4", "leaf", {}),
]
})
mini_ast_graph, _ = generic_ast_graphs.ast_to_graph(
mini_ast_node, mini_ast_spec)
# Allowing 10 children.
nth_child_edges = graph_edge_util.compute_nth_child_edges(
mini_ast_graph, 10)
self.assertEqual(
nth_child_edges,
[("root__root", f"root_children_{i}__child-seq-helper",
f"CHILD_INDEX_{i}") for i in range(5)])
# Allowing 2 children.
nth_child_edges = graph_edge_util.compute_nth_child_edges(
mini_ast_graph, 2)
self.assertEqual(
nth_child_edges,
[("root__root", f"root_children_{i}__child-seq-helper",
f"CHILD_INDEX_{i}") for i in range(2)])