Python decorator with parameters, to run the functions multiple times?

Question

I want to write a python decorator to decorate a test function of a unittest.TestCase, to decide the target host this function should run against. See this example:

class MyTestCase(unittest.TestCase):
    @target_host(["host1.com", "host2.com"])
    def test_my_command(self):
        #do something here against the target host

In the decorated function, I want to be able to execute this test against all hosts, how do I do that? The declaration of target_host is supposed to return a new function, but is it possible to return multiple function that the test runner can execute?

Thanks!

Answer 1

You can return exactly one object (so technically, you could return a collections of functions). If you want to avoid astonishing everyone and if you want to call the result, you better return a single function though. But that function may very well call several other function in a loop... do you see where this leads to?

You need a factory for decorators that return a closure calling the function they're applied to once per set of arguments that factory got. In code (including functools.wraps to keep name and docstring, may be useful or not, I tend to include it by default):

def call_with_each(*arg_tuples):
    def decorate(f):
        @functools.wraps(f)
        def decorator():
            for arg_tuple in arg_tuples:
                f(*arg_tuple)
        return decorator
    return decorate

# useage example:
@call_with_each((3,), (2,)) # note that we pass several singleton tuples
def f(x):
    print x
# calling f() prints "3\n2\n"

Supporting keyword arguments requires more code and perhaps some ugliness, but is possible. If it's always going to be a single argument, the code can be simplified ( def call_with_each(*args) , for arg in args: f(arg) , etc.).

Answer 2

Using a class-based decorator is the way to go if you want your decorator to accept arguments. In my experience they end up being easier to reason about and maintain. They are quite simple, __init__ accepts any arguments for the decorator, __call__ returns the decorated function. One of the gotchas with writing decorators is that they behave completely different if they are passed arguments or not. It's pretty easy to account for this in a class-based decorator, allowing your decorator to accept arguments or not:

class target:
    def __init__(self, targets):
        """Arguments for decorator"""
        self.targets = None
        if not hasattr(targets, '__call__'):
            # check we are actually passed arguments!
            # if targets has __call__ attr, we were called w/o arguments
            self.targets = targets

    def __call__(self, f):
        """Returns decorated function"""
        if self.targets:
            def newf(*args, **kwargs):
                for target in self.targets:
                    f(target)
            return newf
        else:
            return f

Now if we use the decorator with arguments, it will work as expected, calling our function 3 times:

>>> @target([1,2,3])
..: def foo(x): print x
...

>>> foo()
1
2
3

However, if we are not called with arguments, we'll return the original function instead:

>>> @target       
def foo(x): print x
..: 

>>> foo(3)
<<< 3