"""

Implements a function tracer class that can be used to print the call
graph of any Python function.

Lots of test cases are in the test_trace.py file.  They should
demonstrate how to use the classes.  Here's a quick and dirty example:

      >>> t = FunctionTracer()
      >>> t.trace_call(some_func, (arg1, arg2))
      >>> # You can also do this:
      >>> t.start() # starts tracing
      >>> some_func(arg1, arg2)
      >>> t.stop() # stops tracing

WARNING: This code does not deal with exceptions properly and any code
that raises and traps exceptions will produce poorly indented (and
incorrect) results.

Licence: Python License.

Author: Prabhu Ramachandran <prabhu_r at users dot sf dot net>

"""

import sys
import os
import re
import traceback


__MATCH_PATTERN=re.compile(r"^_+")

def get_all_bases(cl, bases):
    for base in cl.__bases__:
        if base not in bases:
            bases.append(base)
    for base in cl.__bases__:
        get_all_bases(base, bases)

def check_if_func_in_class(cl, f_name, co_code):
    my_co_code = 0
    f = 0
    try:
        f = cl.__dict__[f_name]
        my_co_code = f.func_code.co_code
    except KeyError:
        # trying again since funciton could be protected.
        protected_f_name = __MATCH_PATTERN.sub('_', cl.__name__)
        try:
            f = cl.__dict__[protected_f_name+f_name]
            my_co_code = f.func_code.co_code
        except KeyError:
            pass
    if f and (my_co_code == co_code):
        return 1    
    else:
        return 0

def get_class_name (cl, f_name, co_code):
    if check_if_func_in_class(cl, f_name, co_code):
        return cl.__name__
    else:
        bases = []
        get_all_bases(cl, bases)
        for base in bases:
            if check_if_func_in_class(base, f_name, co_code):
                return base.__name__

class FunctionTracer:

    """ This class implements a function tracer for Python.  It can be
    used to trace any Python function.  By trace I mean identify the
    sequence of functions in the call graph. The call sequence is
    printed.  The print function itself is not traced and the current
    implementation simply prints the call graph to sys.stdout.  One
    can subclass this class and use a user defined print statement.

    Examples::

      >>> t = FunctionTracer()
      >>> t.trace_call(some_func, (arg1, arg2))
      >>> t.start() # starts tracing
      >>> some_func(arg1, arg2)
      >>> t.stop() # stops tracing
    
    """

    def __init__(self, verbosity=1, indent="| ", file_info=1, 
                 long_filename=0, output_filename="", dereference_args=0):
        """
        Input Arguments

            verbosity -- Set this to 0 if you only want function
                         entry point information.  Set it to 1
                         (default) if you want a little more verbose
                         output which tells you when the function
                         returns.

            indent --  The text to be used to denote indentation.
                       This makes the output prettier.

            file_info -- If set to non-zero, also prints the
                         filename/line number where function starts.

            long_filename -- If set to 1 it prints the full filename
                             of the file containing the function.

            output_filename -- Outputs the trace to the specified
                               filename.  The contents are appended
                               to the file if the file exists.  The
                               default is to set it to "" and the output
                               will go to sys.stdout.

            dereference_args -- Setting this to non-zero will also print
                                the values of the arguments passed to
                                the function. 

         """        
        self.verbosity = verbosity
        self.indent = indent
        self.file_info = file_info
        self.long_filename = long_filename
        self.out_file = None
        if output_filename:
            self.out_file = open(output_filename, 'a')

        self.dereference_args = dereference_args
        self.n_ind = 0

        # Stores if an exception has been handled once.
        self.hit_exception = 0
        self.last_exception_in = ''
        
        # hack required in some cases. I have no idea why
        self.sys = sys
        self.in_trace = 0
    
    def __del__(self):
        if self.out_file:
            self.out_file.close ()
        self.in_trace = 0
        self.sys.settrace(None)
    
    def trace_call(self, func, args):        
        """ func is the function object to be called, args is a tuple of
        arguments.  This function returns the value of called function."""
        self.sys.settrace(self.trace_dispatch)
        self.in_trace = 1
        ret = apply(func, args)
        self.in_trace = 0
        self.sys.settrace(None)
        return ret
    
    def start(self):
        """ Start tracing. """
        self.sys.settrace(self.trace_dispatch)
        self.in_trace = 1

    def stop(self):
        """ Stop tracing. """
        self.in_trace = 0
        self.sys.settrace(None)

    def set_verbosity(self, verbosity):        
        """ Change the verbosity of the reporting.  Currently this is
        boolean and you can only set verbosity on/off.  A zero value
        is off and non-zero is on."""
        self.verbosity = verbosity

    def set_file_info(self, file_info):
        """If set to non-zero, also prints the filename/line number
        where function starts.  """        
        self.file_info = file_info

    def set_indent(self, indent):
        """ Change the indentation string to be used while tracing."""
        self.indent = indent

    def set_long_filename(self, val):        
        """ If set to 1 then long filenames printed in the call trace
        if set to zero (the default) only the basename is printed."""
        self.long_filename = val

    def set_dereference_args(self, val):        
        """ If set to 1 then the function arguments are dereferenced
        and printed in the call trace."""
        self.dereference_args = val

    def set_output_filename(self, filename):
        """ Change the output file where trace output is dumped."""
        if self.out_file:
            self.out_file.close()
        if filename:
            self.out_file = open(filename, 'a')
        
    def trace_dispatch(self, frame, event, arg):
        """ The main tracer that calls different functions depending
        on the event."""
        if event == 'call':
            if self.hit_exception:
                self.hit_exception = 0
            return self.dispatch_call(frame, arg)
        if event == 'return':
            if self.hit_exception:
                self.hit_exception = 0
            return self.dispatch_return(frame, arg)
        if event == 'exception':
            return self.dispatch_exception(frame, arg)
                
        return self.trace_dispatch

    def _deref_args(self, frame, args):        
        """Given a frame and argument list (of strings), this finds
        the values of the arguments. 'self' is always excluded."""
        arg_values = []
        for a in args:
            if a != 'self':
                try:
                    arg_values.append("%s=%s"%(a, repr(frame.f_locals[a])))
                except (NameError, KeyError, AttributeError):
                    pass                
        return arg_values

    def dispatch_call(self, frame, arg):
        """ Called when a function block is entered."""
        if frame.f_code.co_name != '?':
            self.sys.settrace (None)

            indent = self.n_ind*self.indent
            func_name = frame.f_code.co_name
            func_args = frame.f_code.co_varnames[:frame.f_code.co_argcount]
            file_name = frame.f_code.co_filename
            if not self.long_filename:
                try:
                    file_name = os.path.basename (file_name)
                except AttributeError:
                    print "**Bad, somehow the os module is missing!!**"

            line_num = frame.f_code.co_firstlineno
            class_name = ""
            msg = ""
            try:
                cl = frame.f_locals['self'].__class__
                class_name = get_class_name (cl, func_name,
                                             frame.f_code.co_code)
                if self.file_info:
                    msg = "%s%s.%s%s -- File %s, line %d"\
                          %(indent, class_name, func_name, func_args,
                            file_name, line_num)
                else:
                    msg = "%s%s.%s%s"%(indent, class_name, func_name,
                                       func_args)
            except KeyError:
                if self.file_info:
                    msg = "%s%s%s -- File %s, line %d"\
                          %(indent, func_name, func_args,
                            file_name, line_num)
                else:
                    msg = "%s%s%s"%(indent, func_name, func_args)

            if self.dereference_args:
                arg_vals = self._deref_args(frame, func_args)
                msg = "%s\n%s  arguments: %s"%(msg, indent, arg_vals)

            self.n_ind = self.n_ind + 1
            self.print_func (msg)
            self.sys.settrace (self.trace_dispatch)

        return self.trace_dispatch

    def dispatch_return(self, frame, arg):
        """ Called when a function returns."""
        if frame.f_code.co_name != '?':
            self.sys.settrace (None)
            self.n_ind = self.n_ind - 1
            indent = self.n_ind*self.indent
            if self.verbosity:
                func_name = frame.f_code.co_name
                func_args = frame.f_code.co_varnames[:frame.f_code.co_argcount]
                msg = ""
                try:
                    cl = frame.f_locals['self'].__class__
                    class_name = get_class_name (cl, func_name,
                                                 frame.f_code.co_code)
                    msg = "%sDone %s.%s%s."%(indent, class_name,
                                            func_name, func_args)
                except KeyError:
                    msg = "%sDone %s%s."%(indent, func_name, func_args)

                self.print_func (msg)
            if self.in_trace:
                self.sys.settrace (self.trace_dispatch)

        return self.trace_dispatch

    def dispatch_exception(self, frame, (exc, value, tb)):
        if frame.f_code.co_name != '?':
            self.sys.settrace (None)
            if not self.hit_exception:
                self.hit_exception = 1
                self.last_exception_in = frame.f_code.co_name
            else:
                self.n_ind = self.n_ind - 1
                self.hit_exception = 0

            # Print exception information.
            info = traceback.extract_tb (tb)
            filename, lineno, function, text = info[-1] # last line only
            msg = "Exception: %s:%d: %s: %s (in %s)" %\
                  (filename, lineno, exc.__name__, str(value), function)
            self.print_func (msg)

            self.sys.settrace (self.trace_dispatch)
        return self.trace_dispatch
        
    def print_func(self, msg):        
        """ A print function that prints the trace messages.  This
        function can be changed in a subclass to do something
        different.  No call tracing is done when this function is
        called.  Input Arguments: msg -- is a simple string.

        """
        if self.out_file:
            self.out_file.write(msg+'\n')
            self.out_file.flush()
        else:
            print msg


# --------------------------------------------------
# Lots of test cases are in the test_trace.py file.
# --------------------------------------------------