Using decorator from base class both in base class and derived class

Question

I have some python objects with some methods in which i would like to do some check at the beggining, depending of this check, the method's code would run, or an execption would be raised. Instead of replicating the "check" code at the beginning of every method I though of doing a decorator, I also want the decorator to be embedded inside the class itself, since it is closely related to it. So basically:

instead of this

class A(object):

    def a_method(self):
        if self.check_var is True:
            (some_code)
        else:
            raise Exception

I would like to have this

class A(object):

    def decorator(function):
        def function_wrapper(self, *args, **kwargs):
            if self.check_var is True:
                return function(self, *args, **kwargs)
            else:
                raise Exception
        return function_wrapper

    @decorator
    def a_method(self):
        (some_code)

My first question is, am I going about this right? or is there a better way. I have many methods of the A class that need to have this check, so that is why I don't want to replicate the code unnecessarily.

My second question is, if I go about this the way I described, I run into a problem when I want to derive a class from class A and performe the same decorator checks. Again I don't want to replicate the code, so I want to reuse the decorator in the base class A to performe checks in the derived class. I read about turning the decorator into a @classmethod however when I do this I am able to use the decorator in the derived class but not in the base class anymore!

So basically I would like something like this:

class A(object):

    @classmethod #maybe
    def decorator(function):
        def function_wrapper(self, *args, **kwargs):
            if self.check_var is True:
                return function(self, *args, **kwargs)
            else:
                raise Exception
        return function_wrapper

    @decorator
    def a_method(self):
        (some_code)

class B(A):

    @decorator
    def b_method(self):
        (some_code)

Does anybody know of any clean way to do this?

Answer 1

Since you would prefer to put the decorator inside the class (rather than outside both of them as I suggested in a comment), below shows a way to do it. It makes the decorator a staticmethod instead of a classmethod, and requires using it in a slightly unusual manner, but only within the class.

For more information regarding the necessity of using the decorator like this, see my question Calling class staticmethod within the class body?

class A(object):
    @staticmethod
    def decorator(function):
        def function_wrapper(*args, **kwargs):
            print('in function_wrapper')
            return function(*args, **kwargs)
        return function_wrapper

    @decorator.__func__  #### Note unusual decorator usage inside defining class
    def a_method(self):
        print('in a_method')

class B(A):
    @A.decorator  #### Normal decorator usage outside defining class
    def b_method(self):
        print('in b_method')

One way to avoid having to use __func__ and still keep the definition in the first class would be to postpone turning it into a staticmethod until the very end of the class definition:

class A(object):
    def decorator(function):
        def function_wrapper(*args, **kwargs):
            print('in function_wrapper')
            return function(*args, **kwargs)
        return function_wrapper

    @decorator
    def a_method(self):
        print('in a_method')

    decorator = staticmethod(decorator)  #### convert for use outside this class

class B(A):
    @A.decorator
    def b_method(self):
        print('in b_method')

Yet another way to avoid the __func__ is something like this:

class A(object):
    class Check:
        @staticmethod
        def decorator(function):
            def function_wrapper(*args, **kwargs):
                print('in function_wrapper')
                return function(*args, **kwargs)
            return function_wrapper

    @Check.decorator
    def a_method(self):
        print('in a_method')

class B(A):
    Check = A.Check

    @Check.decorator
    def b_method(self):
        print('in b_method')

Which has the additional advantage of making usage of the decorator very uniform.