Created
September 17, 2018 22:47
-
-
Save RobertLucian/20f0f387eed73b1cd26b51cfc4e086db to your computer and use it in GitHub Desktop.
Call graph generator for python scripts (only python 2.x version)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/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 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Generating the Actual Graphs
To generate the call graph run it this way:
Next, to generate a visual representation of this graph in PNG format, run this:
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