Playing with ast, matplotlib and networkx

gistfile1.py

import ast
import pprint

import networkx as nx

import matplotlib.pyplot as plt

class viz_walker(ast.NodeVisitor):
    def __init__(self):
        self.stack = []
        self.graph = nx.Graph()

    def generic_visit(self, stmt):
        node_name = str(stmt)

        parent_name = None

        if self.stack:
            parent_name = self.stack[-1]

        self.stack.append(node_name)

        self.graph.add_node(node_name)

        if parent_name:
            self.graph.add_edge(node_name, parent_name)

        super(self.__class__, self).generic_visit(stmt)

        self.stack.pop()

class mywalker(ast.NodeVisitor):
    def __init__(self):
        self.def_stack = []

    @property
    def in_func_def(self):
        return len(self.def_stack) > 0

    def proceed(self, stmt):
        super(self.__class__, self).generic_visit(stmt)

    def visit_Call(self, stmt):
        print "*************", stmt.func, stmt.lineno, stmt.args, stmt.kwargs
        pprint.pprint(dir(stmt))

    def visit_FunctionDef(self, stmt):
        func_name = stmt.name

        self.def_stack.append(func_name)

        self.proceed(stmt)

        # Pop this function from the stack
        self.def_stack.pop()

    def visit_For(self, stmt):
        print "stmt: '{0}', target: '{1}', lineno: '{2}', body: '{3}'".format(
            str(stmt), stmt.target.id, stmt.lineno, stmt.body)

        self.proceed(stmt)

    def visit_Name(self, stmt):
        if self.in_func_def:
            print "stmt: '{0}', id: '{1}', lineno: '{2}'".format(
                str(stmt), stmt.id, stmt.lineno)

        self.proceed(stmt)

def main():
    filename = 'test.py'

    with open(filename, 'r') as fin:
        src = fin.read()

    node = ast.parse(src)

    mw = viz_walker()

    mw.visit(node)

    nx.draw(mw.graph)

    plt.show()

if __name__ == "__main__":
    main()

ya 3am koosomk