Python 多层 decorator 内获取原始函数参数字典

0. 在 decorator 里获取原始函数的参数值

项目里做了一个通用锁,使用 decorator 来方便的包住某些需要限制并发的函数。因为并发不是函数级别的,而是根据参数来限制,所以需要把参数传到通用锁的 decorator 里,代码大致如下

def lock_decorator(key=None):
    def _lock_func(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            # TODO: get lock_key
            lock_key = kwargs.get(key, '')
            with LockContext(key=lock_key):
                return func(*args, **kwargs)
        return wrapper
    return _lock_func

@lock_decorator(key='uid')
def apply_recharge(uid, amount):
    # ...

考虑到函数调用不一定都是带着参数名的,就是说调用时不一定所有参数都会进 **kwargs,那就需要从 **args 里面按参数名捞参数

怎么能知道原函数的参数名列表,翻各种手册的得知可以用 inspect.getargspec(func) 来搞到,那么上面的 TODO 部分就可以改写如下

            args_name = inspect.getargspec(func)[0]
            key_index = args_name.index(key)
            if len(args) > key_index:
                lock_key = args[key_index]
            else:
                lock_key = kwargs.get(key, '')

自此,一切都很美好

1. 在 decorator 里获取原始函数的调用参数字典

项目里又做了个通用的 Logger,也做成 decorator 往目标函数一套,就可以打印出调用时的入参和结果,大致如下

def log_decorator():
    def _log_func(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            # TODO: get full args
            print('call func [{}] with args [{}] and kwargs [{}]'.format(func.__name__, args, kwargs))
            ret = func(*args, **kwargs)
            print('func [{}] return [{}]'.format(func.__name__, ret))
            return ret
        return wrapper
    return _log_func

@log_decorator()
def apply_recharge(uid, amount):
    # ...

看起来也还好,不过因为函数可能带默认参数,而且也希望看到 **args 到底传到哪个参数上,还是希望把所有参数按 Key-Value 的形式打印出来,跟处理通用锁一样,用 inspect.getargspec(func) 把参数名和默认值都摸出来,再考虑一下可变参数的情况,对上面的 TODO 部分改写如下

            args_name, _, _, func_defaults = inspect.getargspec(func)
            parsed_kwargs = dict()
            # default args
            default_args = dict()
            default_start = len(args_name, func_defaults)
            for idx, d in enumerate(func_defaults):
                default_args[args_name[default_start + idx]] = d
            parsed_kwargs.update(default_args)
            # args with name
            varargs_start = len(args_name)
            for idx, a in enumerate(args[:varargs_start]):
                parsed_kwargs[args_name[idx]] = a
            # varargs
            if len(args) > varargs_start:
                parsed_kwargs['varargs'] = args[varargs_start:]
            # kwargs
            parsed_kwargs.update(kwargs)
            print('call func [{}] with args [{}]'.format(func.__name__, parsed_kwargs))

到这里,还是很美好

2. 多层 decorator 怎么拿到最原始函数的参数表

注意到上面两个例子里,apply_recharge 都只套了一个 decorator,如果两个一起用会发生什么?

根据 PEP318 里对 decorator 的定义

@dec2
@dec1
def func(arg1, arg2, ...):
    pass

等价于

def func(arg1, arg2, ...):
    pass
func = dec2(dec1(func))

这里就出问题了,dec2 拿到的传入函数其实是 dec1 而不是 func。不过在把 lock_decoratorlog_decorator 混用时,不管谁写前面,func.__name__ 都是原始的函数名,说明也还是有神器的地方做了穿透,但是 inspect.getargspec 又拿不到最底层函数的参数表,导致不管谁前谁后,都有问题

注意到每个 decorator 构建的时候都又封了一个 @functools.wraps(func),这个是干嘛的呢?以前都是无脑用,也没想过为啥要包一层这个,去掉会怎样?

去掉这个 @functools.wraps(func) 后,inspect.getargspec 还是一样的只能拿到最近一层的信息,而之前本来可以拿到底层的 func.__name__ 也变成最近一层的函数名了,说明这里做了穿透。那么去看看代码吧

# functools.py

from _functools import partial, reduce

# update_wrapper() and wraps() are tools to help write
# wrapper functions that can handle naive introspection

WRAPPER_ASSIGNMENTS = ('__module__', '__name__', '__doc__')
WRAPPER_UPDATES = ('__dict__',)
def update_wrapper(wrapper,
                   wrapped,
                   assigned = WRAPPER_ASSIGNMENTS,
                   updated = WRAPPER_UPDATES):
    """Update a wrapper function to look like the wrapped function

       wrapper is the function to be updated
       wrapped is the original function
       assigned is a tuple naming the attributes assigned directly
       from the wrapped function to the wrapper function (defaults to
       functools.WRAPPER_ASSIGNMENTS)
       updated is a tuple naming the attributes of the wrapper that
       are updated with the corresponding attribute from the wrapped
       function (defaults to functools.WRAPPER_UPDATES)
    """
    for attr in assigned:
        setattr(wrapper, attr, getattr(wrapped, attr))
    for attr in updated:
        getattr(wrapper, attr).update(getattr(wrapped, attr, {}))
    # Return the wrapper so this can be used as a decorator via partial()
    return wrapper

def wraps(wrapped,
          assigned = WRAPPER_ASSIGNMENTS,
          updated = WRAPPER_UPDATES):
    """Decorator factory to apply update_wrapper() to a wrapper function

       Returns a decorator that invokes update_wrapper() with the decorated
       function as the wrapper argument and the arguments to wraps() as the
       remaining arguments. Default arguments are as for update_wrapper().
       This is a convenience function to simplify applying partial() to
       update_wrapper().
    """
    return partial(update_wrapper, wrapped=wrapped,
                   assigned=assigned, updated=updated)

原来就是这里耍花样了,把底层函数的 ('__module__', '__name__', '__doc__') 都赋给了 decorator 封起来的这一层,欺骗更上层用 __name__ 去判断时就当我是底层

那我也学这个,把 inspect.getargspec 的地方也处理下不就完了,去看看这个地方是怎么拿参数表的

# inspect.py

def getargspec(func):
    """Get the names and default values of a function's arguments.

    A tuple of four things is returned: (args, varargs, varkw, defaults).
    'args' is a list of the argument names (it may contain nested lists).
    'varargs' and 'varkw' are the names of the * and ** arguments or None.
    'defaults' is an n-tuple of the default values of the last n arguments.
    """

    if ismethod(func):
        func = func.im_func
    if not isfunction(func):
        raise TypeError('{!r} is not a Python function'.format(func))
    args, varargs, varkw = getargs(func.func_code)
    return ArgSpec(args, varargs, varkw, func.func_defaults)

看了下用到了 func.func_codefunc.func_defaults,按 Python 官方文档 https://docs.python.org/2/library/inspect.html 的解释,func_code 是运行时的字节码,从这里面捞参数表果然可行,那是不是我把这两个属性也传递上去就行了呢?改用自己的 wraps 如下

WRAPPER_ASSIGNMENTS = functools.WRAPPER_ASSGNMENTS + ('func_code', 'func_defaults')
def my_wraps(wrapped,
             assigned = WRAPPER_ASSIGNMENTS,
             updated = WRAPPER_UPDATES):
    return _functool.partial(functools.update_wrapper, wrapped=wrapped,
                             assigned=assigned, updated=updated)

运行时报错,看了下错误提示,func_code 不可覆盖,这也对,都是运行时的字节码了,这个覆盖掉那包的这层 decorator 到底还有没有自己的逻辑部分

还是自己动手丰衣足食,既然 func_code 不可覆盖,我自己另外弄一个总可以了吧,而且当前需求是拿到参数表和默认参数,那就直接解出来穿透,也懒得最后再解一次。修改 my_wraps 如下

WRAPPER_ASSIGNMENTS = functools.WRAPPER_ASSIGNMENTS + ('__func_args_name__', '__func_default_args__')

def my_wraps(wrapped,
             assigned = WRAPPER_ASSIGNMENTS,
             updated = functools.WRAPPER_UPDATES):
    if getattr(wrapped, '__func_args_name__', None) is None:
        setattr(wrapped, '__func_args_name__', inspect.getargs(wrapped.func_code)[0])
        func_defaults = getattr(wrapped, 'func_defaults') or ()
        default_args = dict()
        default_start = len(wrapped.__func_args_name__) - len(func_defaults)
        for idx, d in enumerate(func_defaults):
            default_args[wrapped.__func_args_name__[default_start + idx]] = d
        setattr(wrapped, '__func_default_args__', default_args)
    return _functools.partial(functools.update_wrapper, wrapped=wrapped,
                              assigned=assigned, updated=updated)

同时在运行时解参数表,也用一个通用函数来实现

def parse_func(func, *args, **kwargs):
    parsed_kwargs = dict()
    # default args
    parsed_kwargs.update(func.__func_default_args__)
    # args with name
    varargs_start = len(func.__func_args_name__)
    for idx, a in enumerate(args[:varargs_start]):
        parsed_kwargs[func.__func_args_name__[idx]] = a
    # varargs
    if len(args) > varargs_start:
        parsed_kwargs['varargs'] = args[varargs_start:]
    # kwargs
    parsed_kwargs.update(kwargs)

    return parsed_kwargs

这样在 lock_decoratorlog_decorator 里,用 my_wraps 来封装处理,同时在里面用 parse_func 来解析参数,就能拿到完整的参数表了

完整的测试代码见 https://gist.github.com/whusnoopy/9081544f7eaf4e9ceeaa9eba46ff28da