tracer.py - Quick and dirty python 'strace'

tracer.py

#!/usr/bin/python
""" tracer.py: Quick and dirty python 'strace'
"""

import os
import os.path
import sys
import linecache
from functools import wraps

class tracer(object):
    """
    Let's say you want to trace a function/method ``foo``::

        ...
        ...
        def foo(args):
            <stuff you are interested in>
        ...
        ...

    You simply add the following::

        from tracer import tracer
        ...
        ...
        @tracer
        def foo(args):
            <stuff you are interested in>
        ...
        ...

    Now, your function is setup for tracing. Note however, by default nothing
    will be traced and the tracer() function will effectively be a noop until
    there is a DEBUG variable set in the processes environment at runtime.

    So, assuming that the function/method ``foo`` is called when the command
    ``fancyapp`` is run::

        $ fancyapp              # will *not* enable tracing
        $ DEBUG=1 fancyapp      # will enable tracing

    caveats:

        * this tool skips over system modules (ie: anything under
          <sys.prefix>/lib/ ) and builtins. This behavior can be changed by
          overriding the is_ignored() method.
        * this tool currently spits its output to stderr, it might be better to
          send output to a log file instead.
    """

    def __init__(self, fn):
        self.indent = ''
        self.fn = fn

    def is_ignored(self, filename):
        """ is_ignored(filename) -> True or False

        Are calls within this filename skipped during a trace ?
        """
        system_path = os.path.dirname( sys.modules['os'].__file__ )

        return True if (                                # skip over
                filename.startswith(system_path) or     # - system modules
                filename.startswith('<')         or     # - builtins like <string>
                __file__.find(filename) != -1           # - /this/ module
                ) else False

    def trace_fn(self, frame, event, arg):
        """ trace_fn(frame, event, arg) -> trace_fn

        The tracing function that'll be set using the ``sys.settrace()`` call.
        """
        filename = frame.f_code.co_filename
        if self.is_ignored(filename):
            return self.trace_fn

        lineno = frame.f_lineno
        src = linecache.getline(filename, lineno).strip()
        filename = filename if len(filename) < 30 else '...' + filename[-27:]
        if event in ('call', 'return'):
            fn = frame.f_code.co_name
            if event == 'call':
                args = ''
                if frame.f_code.co_argcount > 0:
                    args = ', '.join('%s = %s' % (k, repr(v)) for k, v in frame.f_locals.items())
                sys.stderr.write("%30s +%-5s |%s%s(%s)\n" % (filename, lineno, self.indent, fn, args))
                self.indent += '  '
            else:
                self.indent = self.indent[:-2]
                if src.find('return') != -1:
                    sys.stderr.write("%30s +%-5s |%s%s <= %s\n" % (filename, lineno, self.indent, fn, src.replace('return', '')))
                else:
                    sys.stderr.write("%30s +%-5s |%s%s() <= None\n" % (filename, lineno, self.indent, fn))
        else:
            sys.stderr.write("%30s +%-5s |%s%s\n" % (filename, lineno, self.indent, src))

        return self.trace_fn

    def __call__(self, *args, **kwargs):
        if os.environ.get('DEBUG', None):
            @wraps(self.fn)
            def traceit(*args, **kwargs):
                tracer = sys.gettrace()
                try:
                    sys.settrace(self.trace_fn)
                    return self.fn(*args, **kwargs)
                finally:
                    sys.settrace(tracer)

            return traceit(*args, **kwargs)
        else:
            return self.fn(*args, **kwargs)