#! /usr/bin/env python
######################################################################
#  Copyright (c) by Kevin L. Sitze on 2011-01-23.
#  This code may be used pursuant to the MIT License.
######################################################################

# Source:
# http://code.activestate.com/recipes/577551-trace-decorator-for-debugging/

"""This package provides facilities to attach decorators to classes or
modules (possibly recursively).  A tracing decorator is provided for
tracing function and method calls in your applications.

from decorators import *

@trace
def function(...):
    '''Output goes to <TT>logging.get_logger(__module__)</TT>,
    where <TT>__module__</TT> is the name of the module in which
    the function is defined.  If the current module is '__main__'
    then the root logger will be used.
    '''
    pass

@trace(logger)
def function(...):
    '''Output goes to the specified logger instance.
    '''
    pass

@trace("Logger")
def function(...):
    '''Output goes to the named logger instance.
    '''

class ClassName(object):
    '''Output goes to the ClassName logger.  If ClassName is
    not in the '__main__' module then the name of its module
    will be prefixed to the classname (separated by a dot).
    All methods on the class will be traced.
    '''
    __logger__ = loggerOrLoggerName # optional
    __metaclass__ = decorators.TraceMetaClass
    def method(...):
        pass

class PostDecorate(object):
    def method(...):
        pass
# attach a decorator to an existing class.
attach(trace(logger), PostDecorate)

You can also attach a decorator to an existing module.
"""
from __future__ import absolute_import, print_function

import inspect
import logging
import sys
import types
from functools import wraps
from itertools import chain

import six


try:
    import _thread as thread
except ImportError:
    import thread


FunctionTypes = tuple({types.BuiltinFunctionType, types.FunctionType})

MethodTypes = tuple({types.BuiltinMethodType, types.MethodType})


class _C(object):

    @classmethod
    def class_method(cls):
        pass
    classMethodType = type(class_method)

    @staticmethod
    def static_method():
        pass
    staticMethodType = type(static_method)

    @property
    def property_method(self):
        return None


ClassMethodType = _C.classMethodType
StaticMethodType = _C.staticMethodType
PropertyType = type(_C.property_method)

CallableTypes = tuple({
    types.BuiltinFunctionType, types.FunctionType,
    types.BuiltinMethodType, types.MethodType, ClassMethodType
})

__all__ = (
    'ThreadLocal',
    'TraceMetaClass',
    'attach',
    'get_formatter',
    'set_formatter',
    'get_logger_factory',
    'set_logger_factory',
    'trace'
)


######################################################################
#  Utility functions
######################################################################

MAX_SIZE = 320


def chop(value):
    s = repr(value)
    if len(s) > MAX_SIZE:
        return s[:MAX_SIZE] + '...' + s[-1]
    else:
        return s


def loggable(obj):
    """Return "True" if the obj implements the minimum Logger API
    required by the 'trace' decorator.
    """
    if isinstance(obj, logging.Logger):
        return True
    else:
        return (inspect.isclass(obj)
                and inspect.ismethod(getattr(obj, 'debug', None))
                and inspect.ismethod(getattr(obj, 'isEnabledFor', None))
                and inspect.ismethod(getattr(obj, 'setLevel', None)))


######################################################################
#  LoggerFactory Property
######################################################################

_logger_factory = logging


def get_logger_factory():
    """Retrieve the current factory object for creating loggers.
    The default is to use the logging module.
    """
    global _logger_factory
    return _logger_factory


def set_logger_factory(factory):
    """Set a factory object for creating loggers.  This object must
    publish a method or class named 'get_logger' that takes a string
    parameter naming the logger instance to retrieve.  Logger objects
    returned by this factory must, at a minimum, expose the methods
    'isEnabledFor' and 'debug'.
    """
    global _logger_factory
    _logger_factory = factory


######################################################################
#  class PrependLoggerFactory
######################################################################

class PrependLoggerFactory(object):

    """This is a convenience class for creating new loggers for the
    "trace" decorator.  All loggers created via this class will have a
    user specified prefix prepended to the name of the logger to
    instantiate.
    """

    def __init__(self, prefix='trace'):
        """Construct a new "PrependLoggerFactory" instance that
        prepends the value \var{prefix} to the name of each
        logger to be created by this class.
        """
        self.__prefix = prefix.strip('.')

    @property
    def prefix(self):
        """The value to prefix to each logger created by this factory.
        """
        return self.__prefix

    @prefix.setter
    def prefix(self, value):
        self.__prefix = value.strip('.')

    def get_logger(self, name):
        return logging.getLogger('.'.join((self.__prefix, name)))


######################################################################
#  class ThreadLocal
######################################################################

class ThreadLocal(object):

    """Instances of this class provide a thread-local variable.
    """

    def __init__(self):
        self.__lock = thread.allocate_lock()
        self.__vars = dict()
        self.__init = None

    @property
    def value(self):
        with self.__lock:
            try:
                return self.__vars[thread.get_ident()]
            except KeyError:
                return self.__init

    @value.setter
    def value(self, value):
        with self.__lock:
            self.__vars[thread.get_ident()] = value

    @property
    def initial_value(self):
        with self.__lock:
            return self.__init

    @initial_value.setter
    def initial_value(self, value):
        with self.__lock:
            self.__init = value


######################################################################
#  Formatter functions
######################################################################

def _formatter_self(name, value):
    """Format the "self" variable and value on instance methods.
    """
    __mname = value.__module__
    if __mname != '__main__':
        return '%s = <%s.%s object at 0x%x>' \
            % (name, __mname, value.__class__.__name__, id(value))
    else:
        return '%s = <%s object at 0x%x>' \
            % (name, value.__class__.__name__, id(value))


def _formatter_class(name, value):
    """Format the "klass" variable and value on class methods.
    """
    __mname = value.__module__
    if __mname != '__main__':
        return "%s = <type '%s.%s'>" % (name, __mname, value.__name__)
    else:
        return "%s = <type '%s'>" % (name, value.__name__)


def _formatter_named(name, value):
    """Format a named parameter and its value.
    """
    return '%s = %s' % (name, chop(value))


def _formatter_defaults(name, value):
    return '[%s = %s]' % (name, chop(value))


af_self = _formatter_self
af_class = _formatter_class
af_named = _formatter_named
af_default = _formatter_defaults
af_unnamed = chop
af_keyword = _formatter_named


def get_formatter(name):
    """Return the named formatter function.  See the function
    "set_formatter" for details.
    """
    if name in ('self', 'instance', 'this'):
        return af_self
    elif name == 'class':
        return af_class
    elif name in ('named', 'param', 'parameter'):
        return af_named
    elif name in ('default', 'optional'):
        return af_default
#    elif name in ('anonymous', 'arbitrary', 'unnamed'):
#        return af_anonymous
    elif name in ('keyword', 'pair', 'pairs'):
        return af_keyword
    else:
        raise ValueError('unknown trace formatter %r' % name)


def set_formatter(name, func):
    """Replace the formatter function used by the trace decorator to
    handle formatting a specific kind of argument.  There are several
    kinds of arguments that trace discriminates between:

    * instance argument - the object bound to an instance method.
    * class argument - the class object bound to a class method.
    * positional arguments (named) - values bound to distinct names.
    * positional arguments (default) - named positional arguments with
      default values specified in the function declaration.
    * positional arguments (anonymous) - an arbitrary number of values
      that are all bound to the '*' variable.
    * keyword arguments - zero or more name-value pairs that are
      placed in a dictionary and bound to the double-star variable.

    \var{name} - specifies the name of the formatter to be modified.

        * instance argument - "self", "instance" or "this"
        * class argument - "class"
        * named argument - "named", "param" or "parameter"
        * default argument - "default", "optional"
        * anonymous argument - "anonymous", "arbitrary" or "unnamed"
        * keyword argument - "keyword", "pair" or "pairs"

    \var{func} - a function to format an argument.
    * For all but anonymous formatters this function must accept two
      arguments: the variable name and the value to which it is bound.
    * The anonymous formatter function is passed only one argument
      corresponding to an anonymous value.
    * if \var{func} is "None" then the default formatter will be used.
    """
    if name in ('self', 'instance', 'this'):
        global af_self
        af_self = _formatter_self if func is None else func
    elif name == 'class':
        global af_class
        af_class = _formatter_class if func is None else func
    elif name in ('named', 'param', 'parameter'):
        global af_named
        af_named = _formatter_named if func is None else func
    elif name in ('default', 'optional'):
        global af_default
        af_default = _formatter_defaults if func is None else func
    elif name in ('anonymous', 'arbitrary', 'unnamed'):
        global af_anonymous
        af_anonymous = chop if func is None else func
    elif name in ('keyword', 'pair', 'pairs'):
        global af_keyword
        af_keyword = _formatter_named if func is None else func
    else:
        raise ValueError('unknown trace formatter %r' % name)


######################################################################
#  Decorator: trace
######################################################################

__builtins = (
    '__import__(name,globals={},locals={},fromlist=[],level=-1)',
    'abs(number)',
    'all(iterable)',
    'any(iterable)',
    'apply(object,args=[],kwargs={})',
    'bin(number)',
    'callable(object)',
    'chr(i)',
    'cmp(x,y)',
    'coerce(x,y)',
    'compile(source,filename,mode,flags=0,dont_inherit=0)',
    'delattr(object,name)',
    'dir()',
    'divmod(x,y)',
    'eval(source,globals={},locals={})',
    'execfile(filename,globals={},locals={})',
    'filter(function,sequence)',
    'format(value,format_spec="")',
    'getattr(object,name)',
    'globals()',
    'hasattr(object,name)',
    'hash(object)',
    'hex(number)',
    'id(object)',
    'input(prompt=None)',
    'intern(string)',
    'isinstance(object,klass)',
    'issubclass(C,B)',
    'iter(collection)',
    'len(object)',
    'locals()',
    'map(function,sequence)',
    'max(iterable,key=None)',
    'min(iterable,key=None)',
    'next(iterator)',
    'oct(number)',
    'open(name,mode=0666,buffering=True)',
    'ord(c)',
    'pow(x,y,z=None)',
    'range()',
    'raw_input(prompt=None)',
    'reduce(function,sequence)',
    'reload(module)',
    'repr(object)',
    'round(number,ndigits=0)',
    'setattr(object,name,value)',
    'sorted(iterable,cmp=None,key=None,reverse=False)',
    'sum(sequence,start=0)',
    'unichr(i)',
    'vars()',
    'zip(sequence)'
)

__builtin_defaults = {
    '""': "",
    '-1': -1,
    '0': 0,
    '0666': 0o666,
    'False': False,
    'None': None,
    'True': True,
    '[]': list(),
    '{}': dict()
}

__builtin_functions = None


def __lookup_builtin(name):
    """Lookup the parameter name and default parameter values for
    builtin functions.
    """
    global __builtin_functions
    if __builtin_functions is None:
        builtins = dict()
        for proto in __builtins:
            pos = proto.find('(')
            name, params, defaults = proto[:pos], list(), dict()
            for param in proto[pos + 1:-1].split(','):
                pos = param.find('=')
                if not pos < 0:
                    param, value = param[:pos], param[pos + 1:]
                    try:
                        defaults[param] = __builtin_defaults[value]
                    except KeyError:
                        raise ValueError(
                            'builtin function %s: parameter %s: '
                            'unknown default %r' % (name, param, value))
                params.append(param)
            builtins[name] = (params, defaults)
        __builtin_functions = builtins

    try:
        params, defaults = __builtin_functions[name]
    except KeyError:
        params, defaults = tuple(), dict()
        __builtin_functions[name] = (params, defaults)
        print(
            "Warning: builtin function %r is missing prototype" % name,
            file=sys.stderr)
    return len(params), params, defaults


_ = ThreadLocal()
_.initial_value = False


def trace(_name):
    """Function decorator that logs function entry and exit details.

    \var{_name} a string, an instance of logging.Logger or a function.

    Construct a function or method proxy to generate call traces.
    """
    def decorator(_func):
        """This is the actual decorator function that wraps the
        \var{_func} function for detailed logging.
        """

        def positional(name, value):
            """Format one named positional argument.
            """
            if name is __self:
                return af_self(name, value)
            elif name is __klass:
                return af_class(name, value)
            else:
                return af_named(name, value)

        def wrapper(*__argv, **__kwds):
            if not logger.isEnabledFor(logging.DEBUG) or _.value:
                return _func(*__argv, **__kwds)

            try:
                _.value = True
                params = dict(co_defaults)
                params.update(__kwds)
                params.update(zip(co_varnames, __argv))
                if 'raw_password' in params:
                    params['raw_password'] = '<censored>'

                position = [
                    positional(n, params.pop(n))
                    for n in co_varnames[:len(__argv)]
                ]
                defaults = [
                    af_default(n, params.pop(n))
                    for n in co_varnames[len(__argv):]
                ]
                nameless = (
                    af_unnamed(v) for v in __argv[co_argcount:]
                )
                keywords = (
                    af_keyword(n, params[n]) for n in sorted(params.keys())
                )

                params = ', '.join(
                    filter(None, chain(
                        position,
                        defaults,
                        nameless,
                        keywords)))
                # params = params.replace(', [', '[, ').replace('][, ', ', ')

                enter = [pre_enter]
                if params:
                    enter.append(' ')
                    enter.append(params)
                    enter.append(' ')
                enter.append(')')

                leave = [pre_leave]

                try:
                    logger.debug(''.join(enter))
                    try:
                        try:
                            _.value = False
                            result = _func(*__argv, **__kwds)
                        finally:
                            _.value = True
                    except Exception:
                        ex_type, value, traceback = sys.exc_info()
                        leave.append(' => exception thrown\n\traise ')
                        __mname = ex_type.__module__
                        if __mname != '__main__':
                            leave.append(__mname)
                            leave.append('.')
                        leave.append(ex_type.__name__)
                        if value.args:
                            leave.append('(')
                            leave.append(
                                ', '.join(chop(v) for v in value.args))
                            leave.append(')')
                        else:
                            leave.append('()')
                        raise
                    else:
                        if result is not None:
                            leave.append(' => ')
                            leave.append(chop(result))
                finally:
                    logger.debug(''.join(leave))
            finally:
                _.value = False

            return result

        ####
        #  decorator
        ####

        __self = False
        __klass = False

        def nop(x):
            return x

        __rewrap = nop
        if type(_func) in FunctionTypes:
            # functions do not belong to a class.
            __cname = None
        elif type(_func) in MethodTypes:
            # im_self is None for unbound instance methods.
            # Assumption: trace is only called on unbound methods.
            if _func.im_self is not None:
                __rewrap = classmethod
                __cname = _func.im_self.__name__
                __klass = True
            else:
                __cname = _func.im_class.__name__
                __self = True
            _func = _func.im_func
        else:
            # other callables are not supported yet.
            return _func
        __module = _func.__module__
        __fname = _func.__name__

        # Do not wrap initialization and conversion methods.
        if __fname in ('__init__', '__new__', '__repr__', '__str__'):
            return __rewrap(_func)

        # Generate the Fully Qualified Function Name.

        __fqfn = list()
        if __module != '__main__':
            __fqfn.append(__module)
            __fqfn.append('.')
        if __cname is not None:
            __fqfn.append(__cname)
            __fqfn.append('.')
        __fqfn.append(__fname)
        __fqfn = ''.join(__fqfn)

        if type(_name) in CallableTypes:
            logger = get_logger_factory().get_logger(__fqfn)
        elif loggable(_name):
            logger = _name
        elif isinstance(_name, six.string_types):
            logger = get_logger_factory().get_logger(_name)
        else:
            raise ValueError(
                'invalid object %r: must be a function, a method, '
                'a string or an object that implements the Logger API' % _name)

        pre_enter = ['>>> ', __fqfn, '(']
        pre_enter = ''.join(pre_enter)

        pre_leave = ['<<< ', __fqfn]
        pre_leave = ''.join(pre_leave)

        ####
        #  Here we are really mucking around in function internals.
        #  func_code is the low level 'code' instance that describes
        #  the function arguments, variable and other stuff.
        #
        #  func.func_code.co_argcount - number of function arguments.
        #  func.func_code.co_varnames - function variables names, the
        #      first co_argcount values are the argument names.
        #  func.func_defaults - contains default arguments

        try:
            code = _func.__code__
        except AttributeError:
            co_argcount, co_varnames, co_defaults = \
                __lookup_builtin(_func.__name__)
        else:
            co_argcount = code.co_argcount
            co_varnames = code.co_varnames[:co_argcount]
            if _func.__defaults__:
                co_defaults = dict(
                    zip(co_varnames[-len(_func.__defaults__):],
                        _func.__defaults__))
            else:
                co_defaults = dict()
            if __klass:
                __klass = co_varnames[0]
            if __self:
                __self = co_varnames[0]
        return __rewrap(wraps(_func)(wrapper))

    ####
    #  trace
    ####

#    logging.basicConfig(level = logging.DEBUG)
    if type(_name) in CallableTypes:
        return decorator(_name)
    else:
        return decorator


######################################################################
#  attach: apply decorator to a class or module
######################################################################

def attach_to_property(decorator, klass, k, prop_attr, prop_decorator):
    if prop_attr is not None:
        setattr(klass, k, prop_attr)
        value = decorator(getattr(klass, k))
    else:
        value = None
    # Passing "None" to the property decorator causes the new property
    # to assume the original value of the associated attribute.
    return prop_decorator(value)


def attach_to_class(decorator, klass, recursive=True):
    for k, v in six.iteritems(klass.__dict__):
        t = type(v)
        if isinstance(v, types.FunctionType) or t is ClassMethodType:
            setattr(klass, k, decorator(getattr(klass, k)))
        elif t is StaticMethodType:
            setattr(klass, k, staticmethod(decorator(getattr(klass, k))))
        elif t is PropertyType:
            value = getattr(klass, k)
            value = attach_to_property(
                decorator, klass, k, value.fget, value.getter)
            value = attach_to_property(
                decorator, klass, k, value.fset, value.setter)
            value = attach_to_property(
                decorator, klass, k, value.fdel, value.deleter)
            setattr(klass, k, value)
        elif recursive and inspect.isclass(v):
            attach_to_class(decorator, v, recursive)


def attach(decorator, obj, recursive=True):
    """attach(decorator, class_or_module[, recursive = True])

    Utility to attach a \val{decorator} to the \val{obj} instance.

    If \val{obj} is a module, the decorator will be attached to every
    function and class in the module.

    If \val{obj} is a class, the decorator will be attached to every
    method and subclass of the class.

    if \val{recursive} is "True" then subclasses will be decorated.
    """
    if inspect.ismodule(obj):
        for name, fn in inspect.getmembers(obj, inspect.isfunction):
            setattr(obj, name, decorator(fn))
        for name, klass in inspect.getmembers(obj, inspect.isclass):
            attach_to_class(decorator, klass, recursive)
    elif inspect.isclass(obj):
        attach_to_class(decorator, obj, recursive)


######################################################################
#  class TraceMetaClass
######################################################################

class TraceMetaClass(type):

    """Metaclass to automatically attach the 'trace' decorator to all
    methods, static method and class methods of the class.
    """
    def __new__(mcs, class_name, bases, class_dict):
        klass = super(TraceMetaClass, mcs) \
            .__new__(mcs, class_name, bases, class_dict)
        if '__logger__' in class_dict:
            hook = trace(class_dict['__logger__'])
        else:
            hook = trace
        attach_to_class(hook, klass, False)
        return klass


if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)
    logger = logging.root

    class Test(object):
        __logger__ = logger
        __metaclass__ = TraceMetaClass

        @classmethod
        def class_method(cls):
            pass

        @staticmethod
        def static_method():
            pass

        __test = None

        @property
        def test(self):
            return self.__test

        @test.setter
        def test(self, value):
            self.__test = value

        def method(self):
            pass

    Test.class_method()
    Test.static_method()
    test = Test()
    test.test = 1
    assert 1 == test.test
    test.method()

    class Test(object):

        @classmethod
        def class_method(cls):
            pass

        @staticmethod
        def static_method():
            pass

        __test = None

        @property
        def test(self):
            return self.__test

        @test.setter
        def test(self, value):
            self.__test = value

        def method(self):
            pass

        def __str__(self):
            return 'Test(' + str(self.test) + ')'

    attach(trace(logger), Test)
    Test.class_method()
    Test.static_method()
    test = Test()
    test.test = 1
    assert 1 == test.test
    test.method()
    print(str(test))

    @trace(logger)
    def test(x, y, z=True):
        a = x + y
        b = x * y
        if z:
            return a
        else:
            return b

    test(5, 5)
    test(5, 5, False)

    set_logger_factory(PrependLoggerFactory())

    @trace('main')
    def test(x, *argv, **kwds):
        """Simple test
        """
        return x + sum(argv)

    test(5)
    test(5, 5)
    test(5, 5, False)
    test(5, 5, False, name=10)

    test(*range(50))

    assert test.__doc__ == 'Simple test\n        '
    assert test.__name__ == 'test'

    myzip = trace('main')(zip)
    for i, j in myzip(range(5), range(5, 10)):
        print(i, j)
