3 # Prashanth Ellina, 2007
6 # usage: construct_call_graph -i <your_script>
7 # or cat <your_script> | construct_call_graph
10 generates call graph of given python code file
11 in dot format input for graphviz.
14 * statically tried to figure out functions calls
15 * does not understand classes
16 * algorithm is naive and may not statically find
27 except: import sets; set = sets.Set
29 def annotate_ast_list(ast_list):
31 if code in symbol.sym_name: code = symbol.sym_name[code]
32 else: code = token.tok_name[code]
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)
41 def get_atom_name(atom):
43 first_child_code = first_child[0]
44 if first_child_code != token.NAME: return None
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)
54 second_child_code = second_child[0]
55 if second_child_code != symbol.trailer: return None
57 if len(second_child) < 3: return None
58 if second_child[1][0] == token.LPAR and second_child[-1][0] == token.RPAR:
62 def find_fn_call(ast_list, calls):
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)
68 for item in ast_list[1:]:
69 if isinstance(item, list):
70 find_fn_call(item, calls)
72 def process_fn(fn_ast_list, call_graph):
73 dummy, dummy, func_name = fn_ast_list[:3]
74 dummy, func_name = func_name
77 find_fn_call(fn_ast_list, calls)
79 call_graph[func_name] = list(calls)
81 def construct_call_graph(ast_list, call_graph):
83 if code == symbol.funcdef:
84 process_fn(ast_list, call_graph)
86 for item in ast_list[1:]:
87 if isinstance(item, list):
88 construct_call_graph(item, call_graph)
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)
99 construct_call_graph(ast_list, call_graph)
100 #pprint.pprint(call_graph)
104 dot.append("digraph G {")
105 dot.append("rankdir=LR")
106 for from_fn, to_fns in call_graph.iteritems():
108 dot.append('%s;' % from_fn)
111 if to_fn not in call_graph: continue
112 dot.append('%s -> %s;' % (from_fn, to_fn))
115 return '\n'.join(dot)
117 if __name__ == '__main__':
118 oparser = optparse.OptionParser()
120 oparser.add_option('-i', '--input-file', default=None, metavar='FILE', help='python code file to process')
122 options, args = oparser.parse_args()
124 if options.input_file:
125 python_code = open(options.input_file).read()
127 python_code = sys.stdin.read()
129 dot_code = generate_dot_code(python_code)