callgraph

ack.py

import sys
import callgraph

@callgraph.graphme('m', 'n')
def A(m, n):
    if m == 0:
        return n + 1
    if n == 0:
        return A(m - 1, 1)
    n2 = A(m, n - 1)
    return A(m - 1, n2)

print(A(3,2))
callgraph.view()

callgraph.py

import functools
from collections import defaultdict
import graphviz
import sys

to_trace = {}
def graphme(*args):
    def decorator(f):
        to_trace[f.__name__] = args
        return f
    return decorator

nodes, edges = set(), defaultdict(int)
def trace(frame, event, arg):
    f_name = frame.f_code.co_name
    if f_name in to_trace:
        f_args = ', '.join([f'{x}={y}' for x, y in frame.f_locals.items() if x in to_trace[f_name]])
        c_args = ', '.join([f'{x}={y}' for x, y in frame.f_back.f_locals.items() if x in to_trace[f_name]])
        c_name = frame.f_back.f_code.co_name
        f_node = f'{f_name}({f_args})'
        nodes.add(f_node)
        if c_name != '<module>':
            c_node = f'{c_name}({c_args})'
            edges[(c_node, f_node)] += 1
        # print(f_name, c_name)

sys.settrace(trace)

def view():
    dot = graphviz.Digraph()
    for node in nodes:
        dot.node(str(hash(node)), node, shape='box')
    for (src, dst), num in edges.items():
        dot.edge(str(hash(src)), str(hash(dst)), label=str(num))
    dot.view()
    # print(dot.source)