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?

Objectivity

Just define decorator() outside of, and before, both classes. – Houston

I would like to have the decorator inside the class... the stuff the decorator does is very related to the class – Objectivity

put your base class in a file(module) with the decorator, this way it's quite obvious that they belong together and then just do what @Houston said – Acridine

Alternatively, you could skip decorators, and use metaclasses instead. More work up-front, but no work in the subclasses. – Peristalsis

Besides a metaclass as @o11c, suggests, you could use a class decorator which would work for the base class as well as subclass. An possible issue with either technique is that it would then be difficult to selectively assign the decoration to only some arbitrary set of the class's methods. – Houston

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.

Houston

Answer 2

My first question is, am I going about this right?

As martineau said below, the good practice is put classic decorator outside class.

def get_decorator(function, argument):
    def function_wrapper(*args, **kwargs):
        if argument is True:
            return function(*args, **kwargs)
        else:
            raise Exception
    return function_wrapper

class A(object):
    def __init__(self):
        self.check_var = True
        self.a_method = get_decorator(self.a_method, self.check_var)
    def a_method(self):
        (whatever)

class B(A):
    def __init__(self):
        super(B, self).__init__()
        self.b_method = get_decorator(self.b_method, self.check_var)
    def b_method(self):
        (whatever)

Classic decorator is called during class creation time, which is long before an instance is created. Reference

Unni

I think you missed the point that the OP wants to use the decorator in both class A and sublclass B. Calling/using a classmethod decorator inside the class defining it will result in a TypeError: 'classmethod' object is not callable. – Houston

Yep I know, I just think about multiple inheritance, but it seems not a good solution. – Unni

So the second part of your answer would not address the OP's needs. – Houston

Thanks for your suggestion! Putting some description for the second part. – Unni

Your edit is an improvement, but I have thought of an alternative to workaround the TypeError...see my own just-added answer. – Houston

Willingly accept your authority. Learn new things from your answer. Thank you very much. – Unni