Tracing call flows in Python

Python decorators comes handy when you want to intercept a piece of call flow and profiling technique seems just too verbose.
I use this quite often to analyze a python program to understand it better.

Consider the following piece of contrived python code to illustrate this approach of tracing python call flows.

def f():
    f1('some value')


def f1(result): 
    print result
    f2("f1 result")
    

def f2(result): 
    print result
    f3("f2 result")
    fe("f2 result")
    return "f2 result"

def f3(result): 
    print result
    return "f3 result"

def fe(result): 
    print result

f()

Output:

some value
f1 result
f2 result
f2 result

To trace call flows, sys.settrace comes handy. You can pass your own function that can act as python source code debugger which is specific to your task. In my case, this is a function call tracer.
This function acts as a callback function for various events during execution of code. The function is called with the following arguments -frame, event, and arg. frame is the current stack frame. There are many events, I am particularly interested in 'call'. The callback is called with this event when a function call is made. I leveraged this to create a simple call flow. The output can be configured to generate visual graphs. The code below keeps it simple.

import sys
import os
import linecache

trace_depth = 0

def trace(f):
    
    def globaltrace(frame, why, arg):
        global trace_depth
        if why == "call":
            # function call event , extract information
            .....
            pass 
               
            trace_depth = trace_depth + 1
        return None

    def localtrace(frame, why, arg):
        global trace_depth
        if why == "line":
            # line execution event
            pass
        elif why == "return":
            trace_depth = trace_depth - 1
            # function return event
        return localtrace

    def _f(*args, **kwds):
        sys.settrace(globaltrace)
        result = f(*args, **kwds)
        sys.settrace(None)
        return result

    return _f

Here is a way to record function call event details.

if why == "call":
    # Parent frame details
    p_func = frame.f_back.f_code.co_name
    p_file = frame.f_back.f_code.co_filename
    p_lineinfo = frame.f_back.f_lineno
    p_class = ''
    p_module = ''
    if frame.f_back.f_locals.has_key('self'):
        p_class = frame.f_back.f_locals['self'].__class__.__name__
        p_module = frame.f_back.f_locals['self'].__class__.__module__
    
    # Current frame details
    c_func = frame.f_code.co_name
    c_file = frame.f_code.co_filename
    c_lineinfo = frame.f_lineno
    c_class = ''
    c_module = ''
    if frame.f_locals.has_key('self'):
        c_class = frame.f_locals['self'].__class__.__name__
        c_module = frame.f_locals['self'].__class__.__module__
    # Order is Caller -> Callee
    if trace_depth < 10:
        print 'TraceInfo: ', 
              str(p_func).ljust(20), 
              '----->'.rjust(10) ,
              str(c_func).rjust(10)
    

Now, if i want to checkout the call flow for function f1, then all I would need to do is to instrument it.

@trace
def f1(result): 
    print result
    f2("f1 result")

The output is a pretty good result that shows a typical call flow.

TraceInfo:  _f                       ----->         f1
some value
TraceInfo:  f1                       ----->         f2
f1 result
TraceInfo:  f2                       ----->         f3
f2 result
TraceInfo:  f2                       ----->         fe
f2 result

If you want to record line by line execution, publish it for 'line' event.

if why == "line":
    # record the file name and line number of every trace
    filename = frame.f_code.co_filename
    lineno = frame.f_lineno
    bname = os.path.basename(filename)
    print "%s(%d): %s" % (bname, lineno,linecache.getline(filename, lineno)),

output:

callflow.py(79):     print result
some value
callflow.py(80):     f2("f1 result")
callflow.py(84):     print result
f1 result
callflow.py(85):     f3("f2 result")
callflow.py(90):     print result
f2 result
callflow.py(91):     return "f3 result"
callflow.py(86):     fe("f2 result")
callflow.py(94):     print result
f2 result
callflow.py(87):     return "f2 result"

To look at the function return, you will have to work with 'return' event in localtrace.

elif why == "return":
    trace_depth = trace_depth - 1
    # Current frame details
    c_func = frame.f_code.co_name
    c_file = frame.f_code.co_filename
    c_lineinfo = frame.f_lineno
    c_class = ''
    c_module = ''
    if frame.f_locals.has_key('self'):
        c_class = frame.f_locals['self'].__class__.__name__
        c_module = frame.f_locals['self'].__class__.__module__
    print "TraceInfo<---%s:%s:%s:%s:%s" % (c_file, 
                                              c_lineinfo, 
                                              c_module, 
                                              c_class, 
                                              c_func)

output:

some value
f1 result
f2 result
TraceInfo<---callflow.py:91:::f3
f2 result
TraceInfo<---callflow.py:94:::fe
TraceInfo<---callflow.py:87:::f2
TraceInfo<---callflow.py:80:::f1

The above code can be morphed to provide a wide variety of information.