简体   繁体   中英

How to write a decorator to set the first argument?

I wrote a decorator called apply_first that sets the first argument of the decorated function. Unfortunately this decorator has the wrong signature. Any way around this? I usually use decorator to preserve signature, but this time, I'd like to change it.

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

A better solution:

I ended up writing a decorator that fixes decorators that decorate away the first argument

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

Almost all of the code was copied and appropriately modified from Michele Simionato's decorator module ( GitHub ).

With this, the example above can be fixed as:

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

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM