tracer.py

import os
import linecache
import pprint
import sys
from functools import wraps


class IncludeExclude(object):
    def __init__(self, inclusions=None, exclusions=None):
        self._inclusions = inclusions
        self._exclusions = exclusions

    def filter(self, lst):
        # TODO performance? better design?
        unlist = d = None
        is_dict = isinstance(lst, dict)
        if is_dict:
            d = lst.copy()
            lst = lst.keys()
        elif not isinstance(lst, (list, tuple)):
            lst = [lst]
            unlist = True
        if self._inclusions:
            lst = filter(lambda x: x in self._inclusions, lst)
        if self._exclusions:
            lst = filter(lambda x: x not in self._exclusions, lst)
        if is_dict:
            lst = {k: d[k] for k in lst}
        elif unlist:
            lst = lst[0] if lst else None
        return lst


class Tracer(object):
    def __init__(self, file=None, func=None,
                 locals_ie=None, globals_ie=None):
        self._stream = sys.stdout
        self._filename = None
        self._funcname = None
        assert (
            file is None or
            isinstance(file, IncludeExclude)
        ), file
        assert (
            func is None or
            isinstance(func, IncludeExclude)
        ), func
        self._file = file
        self._func = func
        assert (
            locals_ie is None or
            isinstance(locals_ie, IncludeExclude)
        ), locals_ie
        assert (
            globals_ie is None or
            isinstance(globals_ie, IncludeExclude)
        ), globals_ie
        self._locals_ie = locals_ie
        self._globals_ie = globals_ie

    def locals(self, frame):
        # typically filter frame.f_locals
        # and maybe choose how to do that based on frame contents
        # such as file name, function name, line number
        return ([] if self._locals_ie is None
                else self._locals_ie.filter(frame.f_locals))

    def globals(self, frame):
        # typically filter frame.f_globals
        return ([] if self._globals_ie is None
                else self._globals_ie.filter(frame.f_globals))

    def qual(self, filename, funcname):
        return (
            True if self._file is None
            else self._file.filter(filename)
        ) and (
            True if self._func is None
            else self._func.filter(funcname)
        )

    def show_vars(self, legend, vars):
        indent = 6 * " "
        if vars:
            assert isinstance(vars, dict), vars
            print >> self._stream, indent + legend
            print >> self._stream, '\n'.join(
                [indent + x
                 for x in pprint.pformat(vars).split('\n')]
            )

    def trace_lines(self, frame, event, arg):
        if (
            event == 'line' and
            self.qual(
                os.path.realpath(frame.f_code.co_filename),
                frame.f_code.co_name
            )
        ):
            self.show_vars('LOCALS', self.locals(frame))
            self.show_vars('GLOBALS', self.globals(frame))
            co = frame.f_code
            fname = co.co_filename
            line_no = frame.f_lineno
            print >> self._stream, '{0}:{1}   {2}'.format(
                fname,
                line_no,
                linecache.getline(fname, line_no).rstrip()
            )
            return self.trace_lines

    def __call__(self, f):
        self._filename = os.path.realpath(f.func_code.co_filename)
        self._funcname = f.func_name

        @wraps(f)
        def inner(*args, **kw):
            orig = sys.gettrace()
            try:
                sys.settrace(
                    lambda frame, event, arg:
                    self.trace_lines if event == 'call' else None
                )
                return f(*args, **kw)
            finally:
                sys.settrace(orig)
        return inner


class ShowLocals(Tracer):
    def locals(self, frame):
        return frame.f_locals


class _Locals(Tracer):
    def locals(self, frame):
        L = frame.f_locals
        return {k: v for k, v in frame.f_locals.items()
                if k in self._locals.filter()}


class OnlyTheseLocals(Tracer):
    def __init__(self, varnames):
        Tracer.__init__(self, locals_ie=IncludeExclude(inclusions=varnames))


class ExcludeTheseLocals(Tracer):
    def __init__(self, varnames):
        Tracer.__init__(self, locals_ie=IncludeExclude(exclusions=varnames))


class OnlyThisFile(Tracer):
    def qual(self, filename, funcname):
        return os.path.realpath(filename) == self._filename


class OnlyThisFunction(Tracer):
    def qual(self, filename, funcname):
        return funcname == self._funcname


if __name__ == '__main__':
    def locals_qualifier(filename, funcname, vars):
        if os.path.realpath(__file__) == os.path.realpath(filename):
            return {k: v for k, v in vars.items()
                    if k != 'a'}
        else:
            return False

    class OnlyThisFunctionAndTheseLocals(OnlyThisFunction,
                                         OnlyTheseLocals):
        def __init__(self):
            OnlyTheseLocals.__init__(self, ['a'])

    def subroutine(x, y):
        print 'sum of args is', x + y
        return x + y

    def main():
        a = 1
        b = 2
        a += 4
        q = subroutine(a, b)
        c = 3
        print 'x', a + b + c

    for D in [
        Tracer(),
        ShowLocals(),
        OnlyTheseLocals(['a']),
        ExcludeTheseLocals(['a']),
        OnlyThisFunction(),
        OnlyThisFunctionAndTheseLocals()
    ]:
        # these are all function decorators
        foo = D(main)
        print (20 * '=') + ' ' + D.__class__.__name__
        foo()