Search code examples
pythondecorator

Applying a decorator to every method in a class?


I have decorator @login_testuser applied to method test_1():

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

Is there a way I can apply @login_testuser on every method of the class prefixed with "test_"?

In other words, the decorator would apply to test_1(), test_2() methods below, but not on setUp().

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

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

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

Solution

  • In Python 2.6, a class decorator is definitely the way to go. e.g., here's a pretty general one for these kind of tasks:

    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()")
    

    will get you what you desire. In Python 2.5 or worse, the @decallmethods syntax doesn't work for class decoration, but with otherwise exactly the same code you can replace it with the following statement right after the end of the class TestCase statement:

    TestCase = decallmethods(login_testuser)(TestCase)