How to use decorators on overridable class methods

Question

I have a custom class with multiple methods that all return a code. I would like standard logic that checks the returned code against a list of acceptable codes for that method and raises an error if it was not expected.

I thought a good way to achieve this was with a decorator:

from functools import wraps

def expected_codes(codes):
    def decorator(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            code = f(*args, **kwargs)
            if code not in codes:
                raise Exception(f"{code} not allowed!")
            else:
                return code
        return wrapper
    return decorator

then I have a class like so:

class MyClass:
    @expected_codes(["200"])
    def return_200_code(self):
        return "200"

    @expected_codes(["300"])
    def return_300_code(self):
        return "301" # Exception: 301 not allowed!

This works fine, however if I override the base class:

class MyNewClass:
    @expected_codes(["300", "301"])
    def return_300_code(self):
        return super().return_300_code()  # Exception: 301 not allowed!

I would have expected the above overriden method to return correctly instead of raise an Exception because of the overridden decorator.

From what I've gathered through reading, my desired approach won't work because the decorator is being evaluated at class definition- however I'm surprised there's not a way to achieve what I wanted. This is all in the context of a Django application and I thought Djangos method_decorator decorator might have taken care of this for me, but I think I have a fundamental misunderstanding of how that works.

Answer 1

TL;DR

Use the __wrapped__ attribute to ignore the parent's decorator:

class MyNewClass(MyClass):
    @expected_codes(["300", "301"])
    def return_300_code(self):
        return super().return_300_code.__wrapped__(self) # No exception raised

Explanation

The @decorator syntax is equivalent to:

def f():
    pass
f = decorator(f)

Therefore you can stack up decorators:

def decorator(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        print(f"Calling {f.__name__}")
        f(*args, **kwargs)
    return wrapper

@decorator
def f():
    print("Hi!")

@decorator
def g():
    f()  
g()

#Calling g
#Calling f
#Hi!

But if you want to avoid stacking up, the __wrapped__ attribute is your friend:

@decorator
def g():
    f.__wrapped__()
g()

#Calling g
#Hi!

In short, if you call one of the decorated parent's method in a decorated method of the child class, decorators will stack up, not override one another.

So when you call super().return_300_code() you are calling the decorated method of the parent class which doesn't accept 301 as a valid code and will raise its own exception.

If you want to reuse the original parent's method, the one that simply returns 301 without checking, you can use the __wrapped__ attribute which gives access to the original function (before it was decorated):

class MyNewClass(MyClass):
    @expected_codes(["300", "301"])
    def return_300_code(self):
        return super().return_300_code.__wrapped__(self) # No exception raised