cheyp
cheyp

Reputation: 391

Decorator for function with optional arguments

I have a function my_f with many arguments (some of them can be optional), for example,

def my_f(a, b, opt_arg = 3):
    #do some stuff

Now I want to write some decorator for this function that will do some simple stuff depending on the optional argument opt_arg. I do not know the number of arguments of my_f but I know the name of the optional argument. How I can do this? Is is it possible to do in such way that both variants my_f(1,2, opt_arg=3) and my_f(opt_arg =3, 1,2,3,4) would work correct?

Upvotes: 0

Views: 152

Answers (1)

Martijn Pieters
Martijn Pieters

Reputation: 1121654

Use a **kw variable keywords argument in your decorator wrapper and see if opt_arg is present in it:

from functools import wraps

def yourdecorator(f):
    @wraps(f)
    def wrapper(*args, **kw):
        opt_arg = kw.get('opt_arg')
        if opt_arg is not None:
            print "opt_arg was given and set to {!r}".format(opt_arg)
        return f(*args, **kw)
    return wrapper

Demo:

>>> from functools import wraps
>>> def yourdecorator(f):
...     @wraps(f)
...     def wrapper(*args, **kw):
...         opt_arg = kw.get('opt_arg')
...         if opt_arg is not None:
...             print "opt_arg was given and set to {!r}".format(opt_arg)
...         return f(*args, **kw)
...     return wrapper
... 
>>> @yourdecorator
... def my_f(a, b, opt_arg=3):
...     pass
... 
>>> my_f(42, 81)
>>> my_f(42, 81, opt_arg=3.14)
opt_arg was given and set to 3.14

Upvotes: 1

Related Questions