def test_ast_graph_conforms_to_schema(self):
# Some example code using a few different syntactic constructs, to cover
# a large set of nodes in the schema
root = gast.parse(
textwrap.dedent("""\
def foo(n):
if n <= 1:
return 1
else:
return foo(n-1) + foo(n-2)
def bar(m, n) -> int:
x = n
for i in range(m):
if False:
continue
x = x + i
while True:
break
return x
x0 = 1 + 2 - 3 * 4 / 5
x1 = (1 == 2) and (3 < 4) and (5 > 6)
x2 = (7 <= 8) and (9 >= 10) or (11 != 12)
x2 = bar(13, 14 + 15)
"""))
graph, _ = py_ast_graphs.py_ast_to_graph(root)
# Graph should match the schema
schema_util.assert_conforms_to_schema(graph, py_ast_graphs.SCHEMA)
def test_conforms_to_schema(self):
test_schema = {
graph_types.NodeType("a"):
graph_types.NodeSchema(in_edges=[
graph_types.InEdgeType("ai_0"),
graph_types.InEdgeType("ai_1")
],
out_edges=[graph_types.OutEdgeType("ao_0")
]),
graph_types.NodeType("b"):
graph_types.NodeSchema(in_edges=[graph_types.InEdgeType("bi_0")],
out_edges=[
graph_types.OutEdgeType("bo_0"),
graph_types.OutEdgeType("bo_1")
]),
}
# Valid graph
test_graph = {
graph_types.NodeId("A"):
graph_types.GraphNode(
graph_types.NodeType("a"), {
graph_types.OutEdgeType("ao_0"): [
graph_types.InputTaggedNode(
graph_types.NodeId("B"),
graph_types.InEdgeType("bi_0")),
graph_types.InputTaggedNode(
graph_types.NodeId("A"),
graph_types.InEdgeType("ai_1"))
]
}),
graph_types.NodeId("B"):
graph_types.GraphNode(
graph_types.NodeType("b"), {
graph_types.OutEdgeType("bo_0"): [
graph_types.InputTaggedNode(
graph_types.NodeId("A"),
graph_types.InEdgeType("ai_1"))
],
graph_types.OutEdgeType("bo_1"): [
graph_types.InputTaggedNode(
graph_types.NodeId("B"),
graph_types.InEdgeType("bi_0"))
]
})
}
schema_util.assert_conforms_to_schema(test_graph, test_schema)
def test_does_not_conform_to_schema(self, graph, expected_error):
test_schema = {
graph_types.NodeType("a"):
graph_types.NodeSchema(in_edges=[
graph_types.InEdgeType("ai_0"),
graph_types.InEdgeType("ai_1")
],
out_edges=[graph_types.OutEdgeType("ao_0")
]),
graph_types.NodeType("b"):
graph_types.NodeSchema(in_edges=[graph_types.InEdgeType("bi_0")],
out_edges=[
graph_types.OutEdgeType("bo_0"),
graph_types.OutEdgeType("bo_1")
]),
}
# pylint: disable=g-error-prone-assert-raises
with self.assertRaisesRegex(ValueError, expected_error):
schema_util.assert_conforms_to_schema(graph, test_schema)
def test_encode_maze(self):
"""Tests that an encoded maze is correct and matches the schema."""
maze = np.array([
[1, 1, 1, 1, 1],
[1, 0, 1, 1, 1],
[1, 1, 1, 1, 1],
]).astype(bool)
encoded_graph, coordinates = maze_schema.encode_maze(maze)
# Check coordinates.
expected_coords = []
for r in range(3):
for c in range(5):
if (r, c) != (1, 1):
expected_coords.append((r, c))
self.assertEqual(coordinates, expected_coords)
# Check a few nodes.
self.assertEqual(
encoded_graph[graph_types.NodeId("cell_0_0")],
graph_types.GraphNode(
graph_types.NodeType("cell_xRxD"), {
graph_types.OutEdgeType("R_out"): [
graph_types.InputTaggedNode(
graph_types.NodeId("cell_0_1"),
graph_types.InEdgeType("L_in"))
],
graph_types.OutEdgeType("D_out"): [
graph_types.InputTaggedNode(
graph_types.NodeId("cell_1_0"),
graph_types.InEdgeType("U_in"))
],
}))
self.assertEqual(
encoded_graph[graph_types.NodeId("cell_1_4")],
graph_types.GraphNode(
graph_types.NodeType("cell_LxUD"), {
graph_types.OutEdgeType("L_out"): [
graph_types.InputTaggedNode(
graph_types.NodeId("cell_1_3"),
graph_types.InEdgeType("R_in"))
],
graph_types.OutEdgeType("U_out"): [
graph_types.InputTaggedNode(
graph_types.NodeId("cell_0_4"),
graph_types.InEdgeType("D_in"))
],
graph_types.OutEdgeType("D_out"): [
graph_types.InputTaggedNode(
graph_types.NodeId("cell_2_4"),
graph_types.InEdgeType("U_in"))
],
}))
# Check schema validity.
schema_util.assert_conforms_to_schema(encoded_graph,
maze_schema.build_maze_schema(2))