带参数的Python装饰器，运行函数多次？

Question

我想写一个Python装饰器，用来装饰一个unittest.TestCase里的测试函数，以决定这个函数应该在哪个目标主机上运行。看看这个例子：

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

在这个被装饰的函数里，我希望能够让这个测试在所有主机上执行，我该怎么做呢？target_host的声明应该返回一个新函数，但能不能返回多个函数，让测试运行器去执行呢？

谢谢！

Answer 1

如果你想让你的装饰器可以接受参数，使用基于类的装饰器是个不错的选择。根据我的经验，这种方式更容易理解和维护。其实它们很简单，__init__ 方法可以接收装饰器的任何参数，而 __call__ 方法则返回被装饰的函数。写装饰器时一个需要注意的地方是，如果传入参数，装饰器的行为会完全不同。使用基于类的装饰器可以很容易地处理这个问题，让你的装饰器可以选择性地接受参数：

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

现在如果我们使用带参数的装饰器，它会按预期工作，调用我们的函数三次：

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

>>> foo()
1
2
3

但是，如果我们没有传入参数，它就会返回原始的函数：

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

>>> foo(3)
<<< 3

Answer 2

你只能返回一个对象（技术上讲，你可以返回一组函数）。不过，如果你不想让大家感到惊讶，并且想要调用这个结果，最好还是返回一个单独的函数。不过，这个函数可以在一个循环中调用其他多个函数……你明白这意味着什么吗？

你需要一个装饰器工厂，这个工厂会返回一个闭包，这个闭包会在每组参数下调用它所应用的函数。代码示例如下（包括functools.wraps，用来保持函数的名称和文档字符串，这可能有用也可能没用，我通常默认加上）：

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"

支持关键字参数需要更多的代码，可能会有点复杂，但也是可以做到的。如果参数总是只有一个，那么代码可以简化（比如用def call_with_each(*args)，for arg in args: f(arg)，等等）。