如何对类中的每个方法应用装饰器？

Question

我在方法 test_1() 上使用了一个装饰器 @login_testuser：

class TestCase(object):
    @login_testuser
    def test_1(self):
        print "test_1()"

有没有办法让我在这个类中所有以 "test_" 开头的方法上都应用 @login_testuser 呢？

换句话说，这个装饰器会应用到 test_1()、test_2() 这些方法上，但不会应用到 setUp() 方法上。

class TestCase(object):
    def setUp(self):
        pass

    def test_1(self):
        print "test_1()"

    def test_2(self):
        print "test_2()"

Answer 1

你确定把login_testuser的代码放到setUp里面会不好吗？其实setUp就是用来做这个的：它会在每个测试方法之前运行。

Answer 2

当然可以。你需要遍历这个类的所有属性。检查每一个属性，看看它是不是一个方法，并且名字是否以 "test_" 开头。然后，把它替换成你装饰器返回的那个函数。

大概是这样的：

from inspect import ismethod, getmembers
for name, obj in getmembers(TestCase, ismethod):
   if name.startswith("test_"):
       setattr(TestCase, name, login_testuser(obj))

Answer 3

在Python 2.6中，使用类装饰器是个不错的选择。比如，这里有一个比较通用的装饰器，可以用来处理这类任务：

import inspect

def decallmethods(decorator, prefix='test_'):
    def dectheclass(cls):
        for name, m in inspect.getmembers(cls, inspect.isfunction):
            if name.startswith(prefix):
                setattr(cls, name, decorator(m))
        return cls
    return dectheclass


@decallmethods(login_testuser)
class TestCase(object):
    def setUp(self):
        pass

    def test_1(self):
        print("test_1()")

    def test_2(self):
        print("test_2()")

这样就能实现你想要的效果。在Python 2.5或更早的版本中，@decallmethods这种写法不能用来装饰类，但如果你用的代码完全一样，可以在class TestCase语句结束后，直接用下面的语句替代：

TestCase = decallmethods(login_testuser)(TestCase)