construct_call_graph.py

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