1  #!/usr/bin/env python
 2  """This example was kindly provided by https://github.com/oozie.
 3  
 4  This script draws the AST of a Python module as graph with simple points as nodes.
 5  Astmonkey's GraphNodeVisitor is subclassed for custom representation of the AST.
 6  
 7  Usage: python3 edge-graph-node-visitor.py some_file.py
 8  """
 9  import ast
10  import os
11  import sys
12  
13  import pydot
14  
15  from astmonkey import transformers
16  from astmonkey.visitors import GraphNodeVisitor
17  
18  
19  class EdgeGraphNodeVisitor(GraphNodeVisitor):
20      """Simple point-edge-point graphviz representation of the AST."""
21  
22      def __init__(self):
23          super(self.__class__, self).__init__()
24          self.graph.set_node_defaults(shape='point')
25  
26      def _dot_graph_kwargs(self):
27          return {}
28  
29      def _dot_node_kwargs(self, node):
30          return {}
31  
32      def _dot_edge(self, node):
33          return pydot.Edge(id(node.parent), id(node))
34  
35      def _dot_node(self, node):
36          return pydot.Node(id(node), **self._dot_node_kwargs(node))
37  
38  
39  if __name__ == '__main__':
40      filename = sys.argv[1]
41  
42      node = ast.parse(open(filename).read())
43      node = transformers.ParentChildNodeTransformer().visit(node)
44      visitor = EdgeGraphNodeVisitor()
45      visitor.visit(node)
46      visitor.graph.write(filename + '.dot')
47      os.system('sfdp -Tpng -o {} {}'.format(filename + '.png', filename + '.dot'))