construct_call_graph.py

   1 #!/usr/bin/env python
   2
   3 '''
   4 generates call graph of given python code file
   5 in dot format input for graphviz.
   6
   7 limitations:
   8 * statically tried to figure out functions calls
   9 * does not understand classes
  10 * algorithm is naive and may not statically find
  11   all cases
  12 '''
  13
  14 import sys
  15 import parser
  16 import symbol, token
  17 import pprint
  18 import optparse
  19
  20 try: s = set()
  21 except: import sets; set = sets.Set
  22
  23 def annotate_ast_list(ast_list):
  24     code = ast_list[0]
  25     if code in symbol.sym_name: code = symbol.sym_name[code]
  26     else: code = token.tok_name[code]
  27     ast_list[0] = code
  28
  29     for index, item in enumerate(ast_list):
  30         if index == 0: continue
  31         if isinstance(item, list):
  32             ast_list[index] = annotate_ast_list(item)
  33     return ast_list
  34
  35 def get_atom_name(atom):
  36     first_child = atom[1]
  37     first_child_code = first_child[0]
  38     if first_child_code != token.NAME: return None
  39     return first_child[1]
  40
  41 def get_fn_call_data(ast_list):
  42     if len(ast_list) < 3: return None
  43     first_child, second_child = ast_list[1:3]
  44     first_child_code = first_child[0]
  45     if first_child_code != symbol.atom: return None
  46     fn_name = get_atom_name(first_child)
  47
  48     second_child_code = second_child[0]
  49     if second_child_code != symbol.trailer: return None
  50
  51     if len(second_child) < 3: return None
  52     if second_child[1][0] == token.LPAR and second_child[-1][0] == token.RPAR:
  53         return fn_name
  54     else: return None
  55
  56 def find_fn_call(ast_list, calls):
  57     code = ast_list[0]
  58     if code == symbol.power:
  59         fn_name = get_fn_call_data(ast_list)
  60         if fn_name != None and getattr(__builtins__, fn_name, None) == None: calls.add(fn_name)
  61
  62     for item in ast_list[1:]:
  63         if isinstance(item, list):
  64             find_fn_call(item, calls)
  65
  66 def process_fn(fn_ast_list, call_graph):
  67     dummy, dummy, func_name = fn_ast_list[:3]
  68     dummy, func_name = func_name
  69
  70     calls = set()
  71     find_fn_call(fn_ast_list, calls)
  72
  73     call_graph[func_name] = list(calls)
  74
  75 def construct_call_graph(ast_list, call_graph):
  76     code = ast_list[0]
  77     if code == symbol.funcdef:
  78         process_fn(ast_list, call_graph)
  79
  80     for item in ast_list[1:]:
  81         if isinstance(item, list):
  82             construct_call_graph(item, call_graph)
  83
  84     return call_graph
  85
  86 def generate_dot_code(python_code):
  87     ast = parser.suite(python_code)
  88     ast_list = parser.ast2list(ast)
  89     #annotated_ast_list = annotate_ast_list(ast_list)
  90     #pprint.pprint(annotated_ast_list)
  91
  92     call_graph = {}
  93     construct_call_graph(ast_list, call_graph)
  94     #pprint.pprint(call_graph)
  95
  96     dot = []
  97
  98     dot.append("digraph G {")
  99     dot.append("rankdir=LR")
 100     for from_fn, to_fns in call_graph.iteritems():
 101         if not to_fns:
 102             dot.append('%s;' % from_fn)
 103
 104         for to_fn in to_fns:
 105             if to_fn not in call_graph: continue
 106             dot.append('%s -> %s;' % (from_fn, to_fn))
 107     dot.append("}")
 108
 109     return '\n'.join(dot)
 110
 111 if __name__ == '__main__':
 112     oparser = optparse.OptionParser()
 113
 114     oparser.add_option('-i', '--input-file', default=None, metavar='FILE', help='python code file to process')
 115
 116     options, args = oparser.parse_args()
 117
 118     if options.input_file:
 119         python_code = open(options.input_file).read()
 120     else:
 121         python_code = sys.stdin.read()
 122
 123     dot_code = generate_dot_code(python_code)
 124     print dot_code