Today I’ve been using Python decorators to factor out common functionality in test cases. I ran into the slightly (!) interesting problem of how to define a @decorator(like_this). That is, a decorator that takes a parameter.
It’s fairly trivial to define a decorator that doesn’t take a parameter:
from functools import wraps
def decorator(f):
@wraps(f)
def dec(*args, **kwargs):
# do something useful
return f(*args, **kwargs)
return dec
@decorator
def hello(string):
print string
Okay, it doesn’t do much! @wraps(f) from functools does the magic. The end effect is that the function hello ends up pointing to the decfunction that is dynamically created when the file or module is loaded. That is, @decorator runs when the file is loaded, creates a dec instance and then @wraps ‘renames’ things so that hello points to dec and the f in dec points to the original hello definition (the one with the print string)!
But, what if we want to have a decorator that takes a parameter. i.e. we want to do this:
@pdecorator('hello')
def hello2(string, string2='optional'):
"""
hello2 doc string.
"""
print string, string2
So, as usual, one goes off and does some research. My answer was found here in a very useful post by Elf Sternberg. The example on the linked page has a bit more detail than I needed to understand so I simplified it as:
def pdecorator(name):
def inner(f):
def wrapped(*args, **kwargs):
return f(name, *args, **kwargs)
return wraps(f)(wrapped)
return inner
@pdecorator('hello')
def hello2(string, string2='optional'):
"""
My doc string.
"""
print string, string2
So, what’s going on here? PEP-318 explains what’s going on. Essentially, the syntax of a decorator with arguments like this:
@decorator(arg1, arg2, arg3)
def func(*args):
pass
is actually transformed into:
func = decorator(arg1, arg2, arg3)(func)
That is, decorator has to return a function that can decorate func. So in the above example, wrapped is the function we want to wrap around f. So what does
return wraps(f)(wrapped)
actually do? For that we have to look at wraps from functools.
wraps(f) returns a functools.partial object. A partial object is sort of like a function that is partially called, i.e. it isn’t ‘finished’ yet. e.g. Take the complete function f(a,b) -> a+b. f(2,3) -> 5. However, now imagine that g(f,2) returns a function h(b) which takes a single parameter b such that h(3) -> 5. Effectively h(b) (as we designed it) has partially completed or frozen f() such that the first parameter a is now 2 in the add. In Python it looks like this:
from functools import partial add2 = partial(lambda x,y: x+y, 2) add2(3) 5
i.e. the 2 was captured in the in the x position in the function. Cool, huh!
Anyway, what wraps(...) returns is a partially applied functools.update_wrapper(...) function which has captured the __name__, __doc__ and __dict__ of the function passed as an argument. i.e. wraps(f) returns a partially applied update_wrapper function with f.__name__, f.__doc__ and f.__dict__ frozen into it. However, the wrap hasn’t yet been applied to a function. In our example above, wraps(f)(wrapped) thus then wraps f.__name__, f.__doc__ and f.__dict__ on to the wrapped function.
So coming back to the example (repeated again to stop you scrolling up and down!):
def pdecorator(name):
def inner(f):
def wrapped(*args, **kwargs):
return f(name, *args, **kwargs)
return wraps(f)(wrapped)
return inner
@pdecorator('hello')
def hello2(string, string2='optional'):
"""
My doc string.
"""
print string, string2
inner(f) returns a function which is basically wrapped(f(…)) but has the __name__, __doc__ and __dict__ of f.
i.e. the nice equivalent of
def hello2(...): pass hello2 = wrapped(hello2)
