How to write a decorator to set the first argument?

Question

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

What do you mean *"wrong signature"*? What were you *expecting* the resulting signature to be (if not `*args`, why not)? — jonrsharpe, Jul 30 '15 at 16:37
I wasn't expecting it to be any different to `*args`, but I would like it to be `(y)`. As I mention in my post, the `decorator` module takes care of this for you normally. — , Jul 30 '15 at 16:42
Ah, I'm afraid you can't do that. Note that what you're writing is more-or-less [`functools.partial`](https://docs.python.org/3/library/functools.html#functools.partial), and `functools.wraps` preserves signatures from 3.4 onwards. — jonrsharpe, Jul 30 '15 at 16:42
Or perhaps you can, although it isn't neat or easy - see e.g. http://stackoverflow.com/a/1409496/3001761 — jonrsharpe, Jul 30 '15 at 16:48
Might be able to do it with [`inspect.Signature.bind`](https://docs.python.org/3/library/inspect.html#inspect.Signature.bind) in python 3.4. — NightShadeQueen, Jul 30 '15 at 17:17

score 0 · Accepted Answer · answered Aug 03 '15 at 19:52

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

How to write a decorator to set the first argument?

1 Answers1