How to decorate a parent class and make child classes use it? python

What I want is to create a class decorator to decorate a class and works on subclasses too.

Imagine this class:

class CustomBaseTest(TransactionTestCase):
    def __init__(self, *args, **kwargs):
        ...

    def more_custom_helpers(self):
        ...

and the real Test:

class FooTest(CustomBaseTest):
    def decorate_this_foo_is_ok(self):
        ....

    def decorate_this_fails(self):
        ...

what I want is to use a decorator in CustomBaseTest that finds all methods that starts with ‘decoratte_this_’ and execute custom code after and before. I already have the decorator, something like this:

def class_decorator(klass):
    is_method_test = lambda m: not m.startswith('_') and m.startswith('decorate_this_') and isinstance(getattr(klass, m), MethodType)
    test_methods = filter(is_method_test, dir(klass))

    for method_name in test_methods:
        class_method = getattr(klass, method_name)

        def helper(mname, method):
            @wraps(method)
            ... some logic here
            retval = method(*a, **kw)
            ... more logic here
            return retval
        return wrapper

        fn = MethodType(helper(method_name, class_method), None, klass)
        setattr(klass, method_name, fn)
    return klass

do you know if is possible to do that? and how?

thanks!!!

Answer

Thanks to @Markku and @BrenBarn.

Here is the solution.

First we have a simple decorator:

from functools import wraps
def my_decorator(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        # do some stuff
        retval = func(*args, **kwargs)
        # do more stuff
        return retval
    return wrapper

And the metaclass:

class ProfileMetaByClass(type):
    def __init__(cls, name, bases, dct):
        for method_name, method in dct.items():
            if method_name.startswith('decorate_this_'):
                setattr(cls, method_name, my_decorator(method))
        type.__init__(cls, name, bases, dct)

And that worked for me!