How can I add keyword arguments to a wrapped function in Python 2.7?

Question

I first want to stress that I have searched both the web generally and the Python documentation + StackOverflow specifically very extensively and did not manage to find an answer to this question. I also want to thank anyone taking the time to read this.

As the title suggests, I am writing a decorator in Python, and I want it to add keyword arguments to the wrapped function (please note: I know how to add arguments to the decorator itself, that's not what I'm asking).

Here is a working example of a piece of code I wrote that does exactly that for Python 3 (specifically Python 3.5). It uses decorator arguments, adds keyword arguments to the wrapped function and also defines and adds a new function to the wrapped function.

from functools import wraps

def my_decorator(decorator_arg1=None, decorator_arg2=False):
    # Inside the wrapper maker

    def _decorator(func):
        # Do Something 1

        @wraps(func)
        def func_wrapper(
                *args,
                new_arg1=False,
                new_arg2=None,
                **kwds):
            # Inside the wrapping function
            # Calling the wrapped function
            if new_arg1:
                return func(*args, **kwds)
            else:
                # do something with new_arg2
                return func(*args, **kwds)

        def added_function():
            print("Do Something 2")

        func_wrapper.added_function = added_function
        return func_wrapper

    return _decorator

Now this decorator can be used in the following manner:

@my_decorator(decorator_arg1=4, decorator_arg2=True)
def foo(a, b):
    print("a={}, b={}".format(a,b))

def bar():
    foo(a=1, b=2, new_arg1=True, new_arg2=7)
    foo.added_function()

Now, while this works for Python 3.5 (and I assume for any 3.x), I have not managed to make it work for Python 2.7. I'm getting a SyntaxError: invalid syntax on the first line that tries to define a new keyword argument for the func_wrapper, meaning the line stating new_arg1=False,, when importing the module containing this code.

Moving the new keywords to the start of the argument list of func_wrapper solves the SyntaxError but seems to screw with the wrapped function's signature; I'm now getting the error TypeError: foo() takes exactly 2 arguments (0 given) when calling foo(1, 2). This error disappears if I assign the arguments explicitly, as in foo(a=1, b=2), but that is obviously not enough - unsurprisingly, my new keyword arguments seem to be "stealing" the first two positional arguments sent to the wrapped function. This is something that did not happen with Python 3.

I would love to get your help on this. Thank you for taking the time to read this.

Shay

Answer 1

If you only ever specify the additional arguments as keywords, you can get them out of the kw dictionary (see below). If you need them as positional AND keyword arguments, then I think you should be able to use inspect.getargspec on the original function, and then process args and kw in func_wrapper.

Code below tested on Ubuntu 14.04 with Python 2.7, 3.4 (both Ubuntu-provided) and 3.5 (from Continuum).

from functools import wraps

def my_decorator(decorator_arg1=None, decorator_arg2=False):
    # Inside the wrapper maker

    def _decorator(func):
        # Do Something 1
        @wraps(func)
        def func_wrapper(
                *args,
                **kwds):
            # new_arg1, new_arg2 *CANNOT* be positional args with this technique
            new_arg1 = kwds.pop('new_arg1',False)
            new_arg2 = kwds.pop('new_arg2',None)
            # Inside the wrapping function
            # Calling the wrapped function
            if new_arg1:
                print("new_arg1 True branch; new_arg2 is {}".format(new_arg2))
                return func(*args, **kwds)
            else:
                print("new_arg1 False branch; new_arg2 is {}".format(new_arg2))
                # do something with new_arg2
                return func(*args, **kwds)

        def added_function():
            # Do Something 2
            print('added_function')

        func_wrapper.added_function = added_function
        return func_wrapper

    return _decorator

@my_decorator(decorator_arg1=4, decorator_arg2=True)
def foo(a, b):
    print("a={}, b={}".format(a,b))

def bar():
    pass
    #foo(1,2,True,7) # won't work
    foo(1, 2, new_arg1=True, new_arg2=7)
    foo(a=3, b=4, new_arg1=False, new_arg2=42)
    foo(new_arg2=-1,b=100,a='AAA')
    foo(b=100,new_arg1=True,a='AAA')
    foo.added_function()

if __name__=='__main__':
    import sys
    sys.stdout.flush()
    bar()

Output is

new_arg1 True branch; new_arg2 is 7
a=1, b=2
new_arg1 False branch; new_arg2 is 42
a=3, b=4
new_arg1 False branch; new_arg2 is -1
a=AAA, b=100
new_arg1 True branch; new_arg2 is None
a=AAA, b=100
added_function

Answer 2

To add arguments to an existing function's signature, while making that function behave like a normal python function (correct help, signature and TypeError raising in case of wrong arguments provided) you can use makefun, I developped it specifically to solve this use case.

In particular makefun provides a replacement for @wraps that has a new_sig argument where you specify the new signature. Here is how your example would write:

try:  # python 3.3+
    from inspect import signature, Parameter
except ImportError:
    from funcsigs import signature, Parameter

from makefun import wraps, add_signature_parameters

def my_decorator(decorator_arg1=None, decorator_arg2=False):
    # Inside the wrapper maker

    def _decorator(func):
        # (1) capture the signature of the function to wrap ...
        func_sig = signature(func)
        # ... and modify it to add new optional parameters 'new_arg1' and 'new_arg2'.
        # (if they are optional that's where you provide their defaults)
        new_arg1 = Parameter('new_arg1', kind=Parameter.POSITIONAL_OR_KEYWORD, default=False)
        new_arg2 = Parameter('new_arg2', kind=Parameter.POSITIONAL_OR_KEYWORD, default=None)
        new_sig = add_signature_parameters(func_sig, last=[new_arg1, new_arg2])

        # (2) create a wrapper with the new signature
        @wraps(func, new_sig=new_sig)
        def func_wrapper(*args, **kwds):
            # Inside the wrapping function

            # Pop the extra args (they will always be there, no need to provide default)
            new_arg1 = kwds.pop('new_arg1')
            new_arg2 = kwds.pop('new_arg2')
            
            # Calling the wrapped function
            if new_arg1:
                print("new_arg1 True branch; new_arg2 is {}".format(new_arg2))
                return func(*args, **kwds)
            else:
                print("new_arg1 False branch; new_arg2 is {}".format(new_arg2))
                # do something with new_arg2
                return func(*args, **kwds)

        # (3) add an attribute to the wrapper
        def added_function():
            # Do Something 2
            print('added_function')

        func_wrapper.added_function = added_function
        return func_wrapper

    return _decorator

@my_decorator(decorator_arg1=4, decorator_arg2=True)
def foo(a, b):
    """This is my foo function"""
    print("a={}, b={}".format(a,b))

foo(1, 2, True, 7)  # works, except if you use kind=Parameter.KEYWORD_ONLY above (py3 only)
foo(1, 2, new_arg1=True, new_arg2=7)
foo(a=3, b=4, new_arg1=False, new_arg2=42)
foo(new_arg2=-1,b=100,a='AAA')
foo(b=100,new_arg1=True,a='AAA')
foo.added_function()

help(foo)

It works as you would expect:

new_arg1 True branch; new_arg2 is 7
a=1, b=2
new_arg1 True branch; new_arg2 is 7
a=1, b=2
new_arg1 False branch; new_arg2 is 42
a=3, b=4
new_arg1 False branch; new_arg2 is -1
a=AAA, b=100
new_arg1 True branch; new_arg2 is None
a=AAA, b=100
added_function
Help on function foo in module <...>:

foo(a, b, new_arg1=False, new_arg2=None)
    This is my foo function

So you can see that the exposed signature is as expected, and your users do not see the internals. Note that you can make the two new arguments "keyword-only" by setting kind=Parameter.KEYWORD_ONLY in the new signature, but as you already know this does not work in python 2.

Finally, you might be interested in making your decorator code more readable and robust to no-parenthesis usages, using decopatch. Among other things it supports a "flat" style that is well suited in your case because it removes one level of nesting:

from decopatch import function_decorator, DECORATED

@function_decorator
def my_decorator(decorator_arg1=None, decorator_arg2=False, func=DECORATED):

    # (1) capture the signature of the function to wrap ...
    func_sig = signature(func)
    # ... 

    # (2) create a wrapper with the new signature
    @wraps(func, new_sig=new_sig)
    def func_wrapper(*args, **kwds):
        # Inside the wrapping function
        ...  

    # (3) add an attribute to the wrapper
    def added_function():
        # Do Something 2
        print('added_function')

    func_wrapper.added_function = added_function
    return func_wrapper

(I'm also the author of this one, and created it because I was tired of the nesting and the no-parenthesis handling)