-
Notifications
You must be signed in to change notification settings - Fork 417
/
genastdot.py
71 lines (57 loc) · 2.12 KB
/
genastdot.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
###############################################################################
# AST visualizer - generates a DOT file for Graphviz. #
# #
# To generate an image from the DOT file run $ dot -Tpng -o ast.png ast.dot #
# #
###############################################################################
import argparse
import textwrap
from spi import Lexer, Parser, NodeVisitor
class ASTVisualizer(NodeVisitor):
def __init__(self, parser):
self.parser = parser
self.ncount = 1
self.dot_header = [textwrap.dedent("""\
digraph astgraph {
node [shape=circle, fontsize=12, fontname="Courier", height=.1];
ranksep=.3;
edge [arrowsize=.5]
""")]
self.dot_body = []
self.dot_footer = ['}']
def visit_Num(self, node):
s = ' node{} [label="{}"]\n'.format(self.ncount, node.token.value)
self.dot_body.append(s)
node._num = self.ncount
self.ncount += 1
def visit_BinOp(self, node):
s = ' node{} [label="{}"]\n'.format(self.ncount, node.op.value)
self.dot_body.append(s)
node._num = self.ncount
self.ncount += 1
self.visit(node.left)
self.visit(node.right)
for child_node in (node.left, node.right):
s = ' node{} -> node{}\n'.format(node._num, child_node._num)
self.dot_body.append(s)
def gendot(self):
tree = self.parser.parse()
self.visit(tree)
return ''.join(self.dot_header + self.dot_body + self.dot_footer)
def main():
argparser = argparse.ArgumentParser(
description='Generate an AST DOT file.'
)
argparser.add_argument(
'text',
help='Arithmetic expression (in quotes): "1 + 2 * 3"'
)
args = argparser.parse_args()
text = args.text
lexer = Lexer(text)
parser = Parser(lexer)
viz = ASTVisualizer(parser)
content = viz.gendot()
print(content)
if __name__ == '__main__':
main()