如何编写装饰器来设置第一个参数?

How to write a decorator to set the first argument?

我编写了一个名为 apply_first 的装饰器,用于设置装饰函数的第一个参数。不幸的是,这个装饰器的签名错误。有办法解决这个问题吗?我通常使用decorator来保留签名,但这次我想更改它。

def apply_first(x):
    def decorate(f):
        def g(*args):
            return f(*((x,) + args))
        return g
    return decorate

@apply_first(5)
def add(x,y):
    return x+y

print(add(3))
# prints 8

更好的解决方案:

我最终写了一个装饰器,它修复了装饰掉第一个参数的装饰器

import decorator
import sys

def wrapper_string(pre_decor):
    argspec = decorator.getfullargspec(pre_decor)
    if len(argspec.args) == 0:
        raise TypeError("Couldn't find a first argument to decorate away.")
    allargs = list(argspec.args[1:])
    allshortargs = list(argspec.args[1:])

    if argspec.varargs:
        allargs.append('*' + argspec.varargs)
        allshortargs.append('*' + argspec.varargs)
    elif argspec.kwonlyargs:
        allargs.append('*')  # single star syntax
    for a in argspec.kwonlyargs:
        allargs.append('%s=None' % a)
        allshortargs.append('%s=%s' % (a, a))
    if argspec.varkw:
        allargs.append('**' + argspec.varkw)
        allshortargs.append('**' + argspec.varkw)

    head = "def " + pre_decor.__name__ + "(" + ', '.join(allargs) + "):"
    body = "    return _decorator_(_func_)("+ ', '.join(allshortargs) +")"
    return head + "\n" + body

def update_signature(pre_decor, post_decor, **kw):
    "Update the signature of post_decor with the data in pre_decor"
    post_decor.__name__ = pre_decor.__name__
    post_decor.__doc__ = getattr(pre_decor, '__doc__', None)
    post_decor.__dict__ = getattr(pre_decor, '__dict__', {})
    argspec = decorator.getfullargspec(pre_decor)
    if len(argspec.args) == len(argspec.defaults):
        pos = 1
    else:
        pos = 0
    post_decor.__defaults__ = getattr(pre_decor, '__defaults__', ())[pos:]
    post_decor.__kwdefaults__ = getattr(pre_decor, '__kwdefaults__', None)
    #post_decor.__annotations__ = getattr(pre_decor, '__annotations__', None)
    try:
        frame = sys._getframe(3)
    except AttributeError:  # for IronPython and similar implementations
        callermodule = '?'
    else:
        callermodule = frame.f_globals.get('__name__', '?')
    post_decor.__module__ = getattr(pre_decor, '__module__', callermodule)
    post_decor.__dict__.update(kw)

@decorator.decorator
def decorate_first_arg(dec, pre_decor):
    evaldict = pre_decor.__globals__.copy()
    evaldict['_decorator_'] = dec
    evaldict['_func_'] = pre_decor
    wrapper = compile(wrapper_string(pre_decor), "<string>", "single")
    exec(wrapper, evaldict)
    post_decor = evaldict[pre_decor.__name__]
    update(pre_decor, post_decor)
    return post_decor

几乎所有代码都是从 Michele Simionato 的 decorator 模块 (GitHub) 复制并适当修改的。

有了这个,上面的例子就可以固定为:

@decorate_first_arg
def apply_first(x):
    def decorate(f):
        def g(*args):
            return f(*((x,) + args))
        return g
    return decorate