-
Notifications
You must be signed in to change notification settings - Fork 1
/
visualizer.py
127 lines (108 loc) · 3.85 KB
/
visualizer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
from fractions import Fraction
from dataclasses import dataclass
from typing import Optional, NewType
from utils.errors import EndOfStream, EndOfTokens, TokenError
from utils.datatypes import (
Num,
Bool,
Keyword,
Identifier,
Operator,
NumLiteral,
BinOp,
Variable,
Let,
Assign,
If,
BoolLiteral,
UnOp,
ASTSequence,
AST,
Buffer,
ForLoop,
Range,
Declare,
While,
DoWhile,
Print,
)
from core import RuntimeEnvironment
import graphviz as gv
from math import floor
dot = gv.Digraph()
dot.node_attr.update(shape="box", style="rounded")
node_id = "node_{}"
depth = 0
class ASTViz:
"""
An visualizer for Abstract Syntax Trees. Uses graphviz to create a visual representation of the AST.
Recursively traverses the AST with the treebuilder method.
Parameters
----------
depth: int
The depth of the AST. Used to generate unique node ids for graphviz.
code: str
The code that the AST represents. Used as the label for the graph.
"""
def __init__(self, depth: int = 0, code=None):
self.depth = depth
dot.clear(keep_attrs=False)
dot.attr(label=code)
def treebuilder(self, node: AST, depth: int = 0):
"""
Takes an AST and returns a treebuilder for that AST.
"""
id = node_id.format(self.depth)
self.depth += 1
if type(node) == If:
dot.node(id, "If", shape="diamond")
dot.edge(id, self.treebuilder(node.cond, self.depth))
dot.edge(id, self.treebuilder(node.e1, self.depth))
dot.edge(id, self.treebuilder(node.e2, self.depth))
if type(node) == While:
dot.node(id, "While", shape="invtriangle")
dot.edge(id, self.treebuilder(node.cond, self.depth))
dot.edge(id, self.treebuilder(node.seq, self.depth))
if type(node) == ForLoop:
dot.node(id, "For", shape="invtriangle")
dot.edge(id, self.treebuilder(node.var, self.depth))
dot.edge(id, self.treebuilder(node.val_list, self.depth))
dot.edge(id, self.treebuilder(node.stat, self.depth))
if type(node) == Print:
dot.node(id, "Print")
dot.edge(id, self.treebuilder(node.value, self.depth))
if type(node) == ASTSequence:
dot.node(id, "Sequence", shape="square")
AST_id = "AST_{}_{}"
past_node = id
for _ in range(0, len(node.seq)):
current_node = AST_id.format(_, self.depth)
dot.node(current_node, f"{_}", shape="square", color="grey")
dot.edge(current_node, self.treebuilder(node.seq[_], current_node))
dot.edge(past_node, current_node)
past_node = current_node
if type(node) == Range:
dot.node(id, f"{floor(node.start.value)}")
range_id = "range_{}_{}"
past_node = id
for _ in range(floor(node.start.value) + 1, floor(node.end.value)):
current_node = range_id.format(_, self.depth)
dot.node(current_node, f"{_}")
dot.edge(past_node, current_node)
past_node = current_node
if type(node) == Declare:
dot.node(id, "Declare")
dot.edge(id, self.treebuilder(node.var, self.depth))
dot.edge(id, self.treebuilder(node.value, self.depth))
if type(node) == Variable:
dot.node(id, node.name)
if type(node) == NumLiteral:
dot.node(id, str(node.value))
if type(node) == BinOp:
dot.node(id, node.operator)
dot.edge(id, self.treebuilder(node.left, self.depth))
dot.edge(id, self.treebuilder(node.right, self.depth))
if depth == 0:
dot.format = "png"
dot.render("AST", view=True)
return id