Is there a way to apply a decorator to all classmethods in a class?

44 views Asked by At

I'm not asking this to solve my current problem. (I have an alternate viable plan to achieve my goal.) I'm asking this because I'm curious, and I may learn something from the answer(s). But let me describe the original problem for context:

I have a decorator that I apply to functions in order to wrap them in a custom context manager I wrote. It's pretty simple:

@classmethod
def no_autoupdates(cls):
    def decorator(fn):
        def wrapper(*args, **kwargs):
            coordinator = MaintainedModelCoordinator(auto_update_mode="disabled")
            with cls.custom_coordinator(coordinator):
                fn(*args, **kwargs)
        return wrapper
    return decorator

As an aside (although it's not important to the question), the context manager makes the codebase faster because it "skips all auto-updates", and I decided I could save a lot of time in testing by applying this decorator too all of the tests (that aren't testing the auto-update functionality).

We currently have over 500 tests, and I thought it would be cool if I could just apply the decorator to entire TestCase classes. I found this stack post that demonstrates some tricks you can use to apply a decorator to every method in a class. I initially had some exceptions that I worked around by selectively applying the decorator to methods whose names start with "test_" or "setUp". That seemed to work. Debug prints showed the decorators were applied and there were no errors. I did it by splitting up the decorator method into one that evaluates whether it's being applied to a function or a class, and one (a helper) that is just the bare decorator:

@classmethod
def no_autoupdates_function_decorator(cls, fn):
    disabled_coordinator = MaintainedModelCoordinator(auto_update_mode="disabled")
    def wrapper(*args, **kwargs):
        with cls.custom_coordinator(disabled_coordinator):
            fn(*args, **kwargs)
    return wrapper

@classmethod
def no_autoupdates(cls):
    def decorator(fn_or_class):
        if type(fn_or_class).__name__ == "function":
            # Decorate the supplied method
            return cls.no_autoupdates_function_decorator(fn_or_class)
        else:  # Assume class
            # Decorate every method in the class
            for name, fn in inspect.getmembers(fn_or_class, inspect.isfunction):
                if name.startswith("test_") or name.startswith("setUp"):
                    setattr(fn_or_class, name, cls.no_autoupdates_function_decorator(fn))
            return fn_or_class
    return decorator

The problem was that, although I got debug prints showing that the methods were decorated, when they ran, it wasn't working: the auto-updates were happening and the debug prints showed that the context wasn't getting applied - as if the methods weren't decorated at all. (Note, most of these tests primarily need the context in setUpTestData methods, which are classmethods. For all I know, the instance methods were working.)

As a sanity check, I tried decorating one setUpTestData method directly:

@MaintainedModel.no_autoupdates()
@classmethod
def setUpTestData(cls):
    ...

That didn't work either - same result, but changing the decorator order did work:

@classmethod
@MaintainedModel.no_autoupdates()
def setUpTestData(cls):
    ...

The context was correctly applied. I learned why that worked, but briefly, the @classmethod decorator doesn't return a function (it returns a descriptor object), so a function decorator above it simply doesn't work.

Which leads me, finally, to my question...

Is there a way to decorate classmethods of a class (in an automatic fashion like I used above) that effectively inserts the function decorator between the @classmethod decorator and the function definition?

0

There are 0 answers