CODEWITHSUNDEEP
pythoninheritancedecoratorpython-decorators
wim
wim

Reputation: 362687

Propagating class decorators to inherited classes

import inspect
import functools

def for_all_test_methods(decorator):
    def decorate(cls):
        for name, value in inspect.getmembers(cls, inspect.isroutine):
            if name.startswith('test'):
                setattr(cls, name, test_decorator(getattr(cls, name)))
        return cls
    return decorate

def test_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print(func.__name__, args, kwargs)
        res = func(*args, **kwargs)
        return res
    return wrapper

@for_all_test_methods(test_decorator)
class Potato(object):
    def test_method(self):
        print('in method')

class Spud(Potato):
    def test_derived(self):
        print('in derived')

Now if I create a spud instance the test_method which it has inherited remains decorated, but it has an undecorated method test_derived. Unfortunately, if I add the class decorator onto Spud aswell, then his test_method gets decorated twice!

How do I correctly propagate decorators from the parent class onto the children?

Upvotes: 4

Views: 374

Answers (2)

Andrew Clark
Andrew Clark

Reputation: 208475

Here is how you can accomplish this by using a metaclass instead of decorating the class:

import inspect
import functools

def test_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print(func.__name__, args, kwargs)
        res = func(*args, **kwargs)
        return res
    return wrapper

def make_test_deco_type(decorator):
    class TestDecoType(type):
        def __new__(cls, clsname, bases, dct):
            for name, value in dct.items():
                if name.startswith('test') and inspect.isroutine(value):
                    dct[name] = decorator(value)
            return super().__new__(cls, clsname, bases, dct)
    return TestDecoType

class Potato(object, metaclass=make_test_deco_type(test_decorator)):
    def test_method(self):
        print('in method')

class Spud(Potato):
    def test_derived(self):
        print('in derived')

On Python 2.x you would use __metaclass__ = make_test_deco_type(test_decorator) as the first line of the class body instead of having the metaclass=... portion of the class statement. You would also need to replace super() with super(TestDecoType, cls).

Upvotes: 1

Martijn Pieters
Martijn Pieters

Reputation: 1121834

You cannot avoid decorating derived classes; you can find subclasses of a class after subclasses have been decorated, but not auto-decorate them. Use a metaclass instead of you need that sort of behaviour.

You can do one of two things:

  1. Detect already-decorated methods; if there is a __wrapped__ attribute you have a wrapper:

    def for_all_test_methods(decorator):
    def decorate(cls):
        for name, value in inspect.getmembers(cls, inspect.isroutine):
            if name.startswith('test') and not hasattr(value, '__wrapped__'):
                setattr(cls, name, test_decorator(getattr(cls, name)))
        return cls
    return decorate

  2. Limit the class decorator to direct methods only:

    def for_all_test_methods(decorator):
    def decorate(cls):
        for name, value in cls.__dict__.iteritems():
            if name.startswith('test') and inspect.isroutine(value)):
                setattr(cls, name, test_decorator(getattr(cls, name)))
        return cls
    return decorate

Upvotes: 2

Related Questions