How can I use named arguments in a decorator?

Question

If I have the following function:

def intercept(func):
  # do something here

@intercept(arg1=20)
def whatever(arg1,arg2):
  # do something here

I would like for intercept to fire up only when arg1 is 20. I would like to be able to pass named parameters to the function. How could I accomplish this?

Here's a little code sample :

def intercept(func):
    def intercepting_func(*args,**kargs):
     print "whatever"
     return func(*args,**kargs)
    return intercepting_func

@intercept(a="g")
def test(a,b):
    print "test with %s %s" %(a,b)

test("g","d")

This throws the following exception TypeError: intercept() got an unexpected keyword argument 'a'

Answer 1

from functools import wraps

def intercept(target,**trigger):
    def decorator(func):
        names = getattr(func,'_names',None)
        if names is None:
            code = func.func_code
            names = code.co_varnames[:code.co_argcount]
        @wraps(func)
        def decorated(*args,**kwargs):
            all_args = kwargs.copy()
            for n,v in zip(names,args):
                all_args[n] = v
            for k,v in trigger.iteritems():
                if k in all_args and all_args[k] != v:
                    break
            else:
                return target(all_args)
            return func(*args,**kwargs)
        decorated._names = names
        return decorated
    return decorator

Example:

def interceptor1(kwargs):
    print 'Intercepted by #1!'

def interceptor2(kwargs):
    print 'Intercepted by #2!'

def interceptor3(kwargs):
    print 'Intercepted by #3!'

@intercept(interceptor1,arg1=20,arg2=5) # if arg1 == 20 and arg2 == 5
@intercept(interceptor2,arg1=20)        # elif arg1 == 20
@intercept(interceptor3,arg2=5)         # elif arg2 == 5
def foo(arg1,arg2):
    return arg1+arg2

>>> foo(3,4)
7
>>> foo(20,4)
Intercepted by #2!
>>> foo(3,5)
Intercepted by #3!
>>> foo(20,5)
Intercepted by #1!
>>>

functools.wraps does what the "simple decorator" on the wiki does; Updates __doc__, __name__ and other attribute of the decorator.

Answer 2

You can do this by using *args and **kwargs in the decorator:

def intercept(func, *dargs, **dkwargs):
    def intercepting_func(*args, **kwargs):
        if (<some condition on dargs, dkwargs, args and kwargs>):
            print 'I intercepted you.'
        return func(*args, **kwargs)
    return intercepting_func

It's up to you how you wish to pass in arguments to control the behavior of the decorator.

To make this as transparent as possible to the end user, you can use the "simple decorator" on the Python wiki or Michele Simionato's "decorator decorator"

Answer 3

Remember that

@foo
def bar():
    pass

is equivalent to:

def bar():
    pass
bar = foo(bar)

so if you do:

@foo(x=3)
def bar():
    pass

that's equivalent to:

def bar():
    pass
bar = foo(x=3)(bar)

so your decorator needs to look something like this:

def foo(x=1):
    def wrap(f):
        def f_foo(*args, **kw):
            # do something to f
            return f(*args, **kw)
        return f_foo
    return wrap

In other words, def wrap(f) is really the decorator, and foo(x=3) is a function call that returns a decorator.