Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Call graph generator for python scripts (only python 2.x version)
#!/usr/bin/env python
'''
generates call graph of given python code file
in dot format input for graphviz.
limitations:
* statically tried to figure out functions calls
* does not understand classes
* algorithm is naive and may not statically find
all cases
'''
import sys
import parser
import symbol, token
import pprint
import optparse
try: s = set()
except: import sets; set = sets.Set
def annotate_ast_list(ast_list):
code = ast_list[0]
if code in symbol.sym_name: code = symbol.sym_name[code]
else: code = token.tok_name[code]
ast_list[0] = code
for index, item in enumerate(ast_list):
if index == 0: continue
if isinstance(item, list):
ast_list[index] = annotate_ast_list(item)
return ast_list
def get_atom_name(atom):
first_child = atom[1]
first_child_code = first_child[0]
if first_child_code != token.NAME: return None
return first_child[1]
def get_fn_call_data(ast_list):
if len(ast_list) < 3: return None
first_child, second_child = ast_list[1:3]
first_child_code = first_child[0]
if first_child_code != symbol.atom: return None
fn_name = get_atom_name(first_child)
second_child_code = second_child[0]
if second_child_code != symbol.trailer: return None
if len(second_child) < 3: return None
if second_child[1][0] == token.LPAR and second_child[-1][0] == token.RPAR:
return fn_name
else: return None
def find_fn_call(ast_list, calls):
code = ast_list[0]
if code == symbol.power:
fn_name = get_fn_call_data(ast_list)
if fn_name != None and getattr(__builtins__, fn_name, None) == None: calls.add(fn_name)
for item in ast_list[1:]:
if isinstance(item, list):
find_fn_call(item, calls)
def process_fn(fn_ast_list, call_graph):
dummy, dummy, func_name = fn_ast_list[:3]
dummy, func_name = func_name
calls = set()
find_fn_call(fn_ast_list, calls)
call_graph[func_name] = list(calls)
def construct_call_graph(ast_list, call_graph):
code = ast_list[0]
if code == symbol.funcdef:
process_fn(ast_list, call_graph)
for item in ast_list[1:]:
if isinstance(item, list):
construct_call_graph(item, call_graph)
return call_graph
def generate_dot_code(python_code):
ast = parser.suite(python_code)
ast_list = parser.ast2list(ast)
#annotated_ast_list = annotate_ast_list(ast_list)
#pprint.pprint(annotated_ast_list)
call_graph = {}
construct_call_graph(ast_list, call_graph)
#pprint.pprint(call_graph)
dot = []
dot.append("digraph G {")
dot.append("rankdir=LR")
for from_fn, to_fns in call_graph.iteritems():
if not to_fns:
dot.append('%s;' % from_fn)
for to_fn in to_fns:
if to_fn not in call_graph: continue
dot.append('%s -> %s;' % (from_fn, to_fn))
dot.append("}")
return '\n'.join(dot)
if __name__ == '__main__':
oparser = optparse.OptionParser()
oparser.add_option('-i', '--input-file', default=None, metavar='FILE', help='python code file to process')
options, args = oparser.parse_args()
if options.input_file:
python_code = open(options.input_file).read()
else:
python_code = sys.stdin.read()
dot_code = generate_dot_code(python_code)
print dot_code
@RobertLucian

This comment has been minimized.

Copy link
Owner Author

commented Sep 17, 2018

Generating the Actual Graphs

To generate the call graph run it this way:

cat your_module.py | python construct_call_graph.py > your_calls.dot

Next, to generate a visual representation of this graph in PNG format, run this:

cat your_calls.dot | dot -Tpng > your_call_stack.png

Requirements

As for what dependencies are needed, make sure to run this with Python 2.x and have Graphviz installed (that is the dot command). To install Graphviz use:

apt-get update -qq && apt-get install -y --no-install-recommends graphviz

I generally prefer to run this with Docker as that gives me a really clean environment. I do something like this, but everyone can alter it to their own requirements:

docker container run -it --rm -v $PWD:/home -w /home python:2.7-stretch bash

Where $PWD is the path to this gist's file and to the module I want to analyze.

Last Notes

I gotta thank @prashanthellina for providing this simple and yet powerful tool. I'm developing something in micropython and there pretty much isn't a way to profile a script in this language, so I've got to resort to a static analyzer. As much as simple this script is, this thing saved me quite some time. There's also the blog post I've followed:
http://code.prashanthellina.com/code/construct_call_graph.py

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.