Generate a call graph for the enclosed context

graph_stack.py

import contextlib
import os.path
from pycallgraph import PyCallGraph
from pycallgraph.output import GraphvizOutput

@contextlib.contextmanager
def graph_stack(filename):
    'Generate a call graph for the enclosed context'
    go = GraphvizOutput(output_file=filename)
    go.output_type = os.path.splitext(filename)[-1][1:]
    with PyCallGraph(output=go):
        yield

# --------

if __name__ == '__main__':
    def do_hello(times):
        for i in xrange(times):
            hello()

    def hello():
        print('Hello World!')

    with graph_stack('test.pdf'):
        do_hello(5)