This module allows tracing the calls of a running module by using the `sys.settrace()` mechanism.

trace_calls.py

# -*- coding: utf-8 -*-
"""
This module allows tracing the calls of a running module by using
the `sys.settrace()` mechanism.

Usage
-----

The simplest usage is following:

    from top_level_module import trace_calls
    trace_calls.enable_trace_func()
    # ...
    trace_calls.disable_trace_func()

The code between the two calls will be traced.

One can set the maximum depth of the trace and the output amy be customised by
providing a callback using `trace_calls.trace_call_impl = func_name`.
Three callbacks are provided to print only the name, the name and arguments
or the name, the arguments the file name and line number.

To temporarly disable the tracing set `trace_calls.trace_func_enabled = False`
but be aware that enabling it should be in the same function.

To filter out functions by name set the `trace_calls.hide_func_from_trace`
to a list of names. By default the list in `trace_calls.STD_HIDE_FUNCS` is used.

Another predefined filter detects private functions (the ones that start
and end with `__`) if `trace_calls.hide_private = False`.

Created on Wed Mar 25 20:02:12 2015.

@author: Nicu Tofan <[email protected]>
"""
from __future__ import print_function

import os
import sys

STD_HIDE_FUNCS = ['_exitfunc', '_run_exitfuncs', '__contains__', '__get__', '__hash__']

trace_func_enabled = False
trace_func_enabled.__doc__ = """Enable/disable tracking

This variable is set/reset by `enable_trace_func()` and `disable_trace_func()`.
It can be querried to find out the state.
"""

hide_func_from_trace = STD_HIDE_FUNCS
hide_func_from_trace.__doc__ = """List of function names to ignore.

You can replace or append to this list.
"""

hide_private = False
hide_private.__doc__ = """Show/hide private functions.

It can be querried to find out the state.
"""

__current_hide__ = (None, None)
__max_depth__ = sys.maxint

def enable_trace_func():
    """
    Enable tracing.

    Examples
    --------
    >>> from top_level_module import trace_calls
    >>> trace_calls.enable_trace_func()
    >>> # ...
    >>> trace_calls.disable_trace_func()

    """
    global trace_func_enabled
    trace_func_enabled = True
    sys.settrace(__trace__func__)

def disable_trace_func():
    """
    Disable tracing.

    Examples
    --------
    >>> from top_level_module import trace_calls
    >>> trace_calls.enable_trace_func()
    >>> # ...
    >>> trace_calls.disable_trace_func()
    """
    global trace_func_enabled
    trace_func_enabled = False

def set_max_depth(val):
    """
    Set maximum depth for which functions are printed.

    Examples
    --------
    >>> from top_level_module import trace_calls
    >>> trace_calls.set_max_depth(4)
    >>> trace_calls.enable_trace_func()
    >>> # ...
    >>> trace_calls.disable_trace_func()
    """
    global __max_depth__
    if val == None:
        __max_depth__ = sys.maxint
    else:
        __max_depth__ = val*2

def traceprint_name(f_code, indent, arg):
    """
    traceprint implementation that prints the name of the function

    Examples
    --------
    >>> from top_level_module import trace_calls
    >>> trace_calls.trace_call_impl = trace_calls.traceprint_name
    >>> trace_calls.enable_trace_func()
    >>> # ...
    >>> trace_calls.disable_trace_func()
    """
    print('-' * indent[0] + '>',
          f_code.co_name)

def traceprint_name_arg(f_code, indent, arg):
    """
    traceprint implementation that prints the name of the function and arguments

    Examples
    --------
    >>> from top_level_module import trace_calls
    >>> trace_calls.trace_call_impl = trace_calls.traceprint_name_arg
    >>> trace_calls.enable_trace_func()
    >>> # ...
    >>> trace_calls.disable_trace_func()
    """
    print('-' * indent[0] + '>',
          f_code.co_name,
          arg)

def traceprint_name_fname_arg(f_code, indent, arg):
    """
    traceprint implementation that prints the name of the function,
    the arguments and the location (file and line).

    Examples
    --------
    >>> from top_level_module import trace_calls
    >>> trace_calls.trace_call_impl = trace_calls.traceprint_name_fname_arg
    >>> trace_calls.enable_trace_func()
    >>> # ...
    >>> trace_calls.disable_trace_func()
    """
    print('-' * indent[0] + '>',
          f_code.co_name,
          '  ',
          os.path.split(f_code.co_filename)[1],
          '[',
          f_code.co_firstlineno,
          ']',
          arg)

trace_call_impl = traceprint_name_fname_arg
trace_call_impl.__doc__ = """The callback to use on function entry.

You can provide your own function that follows this layout:
   >>> def traceprint_name_fname_arg(f_code, indent, arg):
   >>>     pass
or you can use one of the predefined functions.
"""

def __trace__func__(frame, event, arg, indent=None):
    """
    Actual trace function that is set via `sys.settrace()`
    """
    global trace_func_enabled
    global hide_func_from_trace
    global __current_hide__
    global trace_call_impl
    global __max_depth__
    global hide_private

    if indent is None:
        indent = [0]

    if not trace_func_enabled:
        return None

    if event == "call":
        indent[0] += 2

        if __current_hide__[0] != None:
            return __trace__func__

        fname = frame.f_code.co_name
        if fname in hide_func_from_trace:
            __current_hide__ = (fname, indent)
            return __trace__func__

        if fname.startswith('<') and fname.endswith('>'):
            __current_hide__ = (fname, indent)
            return __trace__func__

        if hide_private:
            if fname.startswith('__') and fname.endswith('__'):
                __current_hide__ = (fname, indent)
                return __trace__func__

        if __max_depth__ < indent[0]:
            __current_hide__ = (fname, indent)
            return __trace__func__

        trace_call_impl(frame.f_code, indent, arg)

    elif event == "return":
        fname = frame.f_code.co_name
        if __current_hide__[0] != None:
            if indent[0] == __current_hide__[1]:
                if fname == __current_hide__[0]:
                    __current_hide__ = (None, None)
        else:
            print('<' + '-' * indent[0], fname)

        indent[0] -= 2

    return __trace__func__