Hooking every function call in Python

Question

I have a large codebase, which consists of thousands of functions.

I want to enable code execution before and after every function call, when the functions starts and when it ends.

Is there a way to do it without recompiling Python, or adding code to each function? Is there a way to hook every function call in my code?

Answer 1

Yes, you can use either the sys.settrace() or sys.setprofile() functions to register a callback and handle 'call' and perhaps 'return' events. However, this can slow down your code considerably. Calling a function has overhead, adding another function call for every function call adds more overhead.

By default, the sys.settrace() hook is only called for calls (where call indicates a new scope being entered, including class bodies and list, dict and set comprehensions, as well as generator expressions), but you can optionally return a trace function to be called for the scope just entered. If you are only interested in calls then just return None from the trace function. Note that this lets you be selective about what scopes you gather more information about. sys.settrace() only reports on Python code, not built-in callables or those defined in compiled extensions.

The sys.setprofile() hook is called for calls to both Python functions and builtins and compiled extension objects, and the same callback is also called whenever a call returns or an exception was raised. Unfortunately it’s not possible to distinguish between a Python function returning None or raising an exception.

In both cases you are given the current frame, as well as the event name and arg, usually set to None but for some events to something more specific:

def call_tracer(frame, event, arg):
    # called for every new scope, event = 'call', arg = None
    # frame is a frame object, not a function!
    print(f"Entering: {frame.f_code.co_name}")
    return None

sys.settrace(call_tracer)

When using sys.settrace() returning a function object instead of None lets you trace other events within the frame, this is the 'local' trace function. You can re-use the same function object for this. This slows things down more, because now you are calling a function for every line of source code. The local trace function is then called for 'line', 'exception' and 'return' events, but you can disable per-line events by setting frame.f_trace_lines = False (requires Python 3.7 or newer).

Here is a short demo of both hooks (assuming Python 3.7 or newer is used); it ignores the exception event option:

import sys

# demo functions, making calls and returning things

def foo(bar, baz):
    return bar(baz)

def spam(name):
    print(f"Hello, {name}")
    return [42 * i for i in range(17)]

# trace functions, one only call events, another combining calls and returns

def call_tracer(frame, event, arg):
    # called for every new scope, event = 'call', arg = None
    # frame is a frame object, not a function!
    print(f"Entering: {frame.f_code.co_name}")
    return None

def call_and_return_tracer(frame, event, arg):
    if event == 'call':
        print(f"Entering: {frame.f_code.co_name}")
        # for this new frame, only trace exceptions and returns
        frame.f_trace_lines = False
        return call_and_return_tracer
    elif event == 'c_call':
        print(f"Entering: {arg.__name__}")
    elif event == 'return':
        print(f"Returning: {arg!r}")
    elif event == 'c_return':
        print(f"Returning from: {arg.__name__}")

if __name__ == '__main__':
    sys.settrace(call_tracer)
    foo(spam, "world")
    print()

    sys.settrace(call_and_return_tracer)
    foo(spam, "universe")
    print()
    sys.settrace(None)

    sys.setprofile(call_and_return_tracer)
    foo(spam, "profile")

Running this outputs:

Entering: foo
Entering: spam
Hello, world
Entering: <listcomp>

Entering: foo
Entering: spam
Hello, universe
Entering: <listcomp>
Returning: [0, 42, 84, 126, 168, 210, 252, 294, 336, 378, 420, 462, 504, 546, 588, 630, 672]
Returning: [0, 42, 84, 126, 168, 210, 252, 294, 336, 378, 420, 462, 504, 546, 588, 630, 672]
Returning: [0, 42, 84, 126, 168, 210, 252, 294, 336, 378, 420, 462, 504, 546, 588, 630, 672]

Entering: foo
Entering: spam
Entering: print
Hello, profile
Returning from: print
Entering: <listcomp>
Returning: [0, 42, 84, 126, 168, 210, 252, 294, 336, 378, 420, 462, 504, 546, 588, 630, 672]
Returning: [0, 42, 84, 126, 168, 210, 252, 294, 336, 378, 420, 462, 504, 546, 588, 630, 672]
Returning: [0, 42, 84, 126, 168, 210, 252, 294, 336, 378, 420, 462, 504, 546, 588, 630, 672]
Returning: None

If possible to alter the code, add decorators only to the functions you want to trace, so you can limit the overhead. You can even automate this if you are prepared to write some code to do the alterations; with the ast module you can parse code into object trees that can be transformed, including adding in @decorator syntax. This isn't that simple but really worth it if your codebase is large. See the Green Tree Snakes project for more in-depth documentation on how to do that.