Liuw’s Thinkpad

我们的项目需要改Xen用户空间的工具。Xen用户空间的工具是用Python写的，代码量还不小，用人肉方法去跟踪执行路径那是很笨的办法。我们需要trace调用，又不能改变函数原来的行为，那么最好的办法就是用decorator了。上网找一下，已经有人做过这样的工作了，我综合一下，改成下面的样子。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Decorators used to trace function/method calls
# derive from:
# http://wordaligned.org/svn/etc/echo/echo.py
# http://svn.osafoundation.org/chandler/trunk/chandler/util/autolog.py
#
# liuw
# 2010-5-27

import types

# Should use python log facility?
def do_log(s):
    print s

def _format_arg_value(arg_val):
    arg, val = arg_val
    return '%s=%r' % (arg, val)

def log_wrapper(fn, display_name=None):
    import functools
    if not display_name: display_name = fn.__name__

    code = fn.func_code
    arg_count = code.co_argcount
    arg_names = code.co_varnames[:arg_count]
    fn_defaults = fn.func_defaults or list()
    arg_defs = dict(zip(arg_names[-len(fn_defaults):], fn_defaults))

    @functools.wraps(fn)
    def wrapped(*v, **k):
        global indent

        positional = map(_format_arg_value, zip(arg_names, v))
        defaulted = [_format_arg_value((a, arg_defs[a]))
                     for a in arg_names[len(v):] if a not in k]
        nameless = map(repr, v[arg_count:])
        keyword = map(_format_arg_value, k.items())

        args = positional + defaulted + nameless + keyword

        do_log('%s(%s)' % (display_name, ', '.join(args)))
        return fn(*v, **k)
    return wrapped

def log_class(cls, methodAsFunction=False):
    """Python does not support class decorator, use it this way:
        class ClassName:
            ...
        ClassName = log_class(ClassName)
    """
    # disallow ALL builtin functions/methods
    allow = lambda s: not (s.startswith('__') and s.endswith('__'))

    names = cls.__dict__.keys()

    for name in names:
        if not allow(name): continue

        value = getattr(cls, name)

        if methodAsFunction and callable(value):
            setattr(cls, name, log_wrapper(value))
        elif isinstance(value, types.MethodType):
            # a normal instance method
            if value.im_self == None:
                setattr(cls, name, log_wrapper(value))

            # class and static method are more complex
            # class method
            elif value.im_self == cls:
                w = log_wrapper(value.im_func,
                               display_name='%s.%s' % (cls.__name__,
                                                       value.__name__))
                setattr(cls, name, classmethod(w))
            else:
                assert False

        # static method
        elif isinstance(value, types.FunctionType):
            w = log_wrapper(value,
                            display_name='%s.%s' % (cls.__name__,
                                                    value.__name__))
            setattr(cls, name, staticmethod(w))
    return cls

def log_module(m):
    # import m if m hasn't been imported
    if isinstance(m, str):
        d = {}
        exec 'import %s' % m in d
        import sys
        m = sys.modules[m]

    names = m.__dict__.keys()
    for name in names:
        value = getattr(m, name)
        if isinstance(value, type):
            setattr(m, name, log_class(value))
        elif isinstance(value, types.FunctionType):
            setattr(m, name, log_wrapper(value))

# ------------------------------TEST----------------------------------
class C1:
    @classmethod
    def f1(self, a, b=1, c=2):
        """f1"""
        pass

    @staticmethod
    def f2(d , e=3, f=4):
        """f2"""
        pass

    def f3(self, g):
        """f3"""
        pass
C1 = log_class(C1)

@log_wrapper
def test(x):
    pass

if __name__ == '__main__':
    inst = C1()
    inst.f1(10)
    inst.f2(20)
    inst.f3(1000)
    test(321)

functools.wraps这个decorator可以保证被修饰函数/方法的其他attribute与被wrap之前无异。比较让我感慨的是，我参考的有一部分是Chandler的代码，这个项目现在仍健在，实在不易。