How can I add keyword arguments to a wrapped function in Python 2.7?
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
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)