Patching decorators in Python

Being a syntactic sugar decorators are essentially a usually callable (e.g. a function or a class) that wraps another callable. And essentially they are applied/called when the module is being read.

For example, the following two chunks of codes are identical:

@decorator
def function():
    return 42

def function():
    return 42

function = decorator(function)

As we can see the decorator in this example is called when the file/module containing the above would be read. So how do we patch such an object?

Obviously, it won't be as straight forward as patching a standalone function as patching the decorator would involve playing with the module import mechanisms a bit. Here's the small recipe I use when I need to patch a decorator:

import importlib
import sys

from dataclasses import dataclass
from types import ModuleType, TracebackType
from typing import Callable, Optional, Type


@dataclass
class PatchDecorator:

    source_module_name: str  # where the decorator is defined
    destination_module_name: str  # where the decorator is imported/used
    decorator_name: str  # name of the decorator
    replacement_decorator_obj: Callable  # callable to replace the decorator

    def setup(self) -> ModuleType:

        sys.modules.pop(self.source_module_name, None)
        sys.modules.pop(self.destination_module_name, None)

        source_module = importlib.import_module(self.source_module_name)
        setattr(source_module, self.decorator_name, self.replacement_decorator_obj)

        patched_destination_module = importlib.import_module(
            self.destination_module_name
        )
        return patched_destination_module

    def cleanup(self) -> None:

        sys.modules.pop(self.source_module_name, None)
        sys.modules.pop(self.destination_module_name, None)

    def __enter__(self) -> ModuleType:
        return self.setup()

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_value: Optional[BaseException],
        traceback: TracebackType,
    ) -> None:
        self.cleanup()

Usage examples:

The following examples assume:

• The decorator is named decorator and is defined in module utils.py
• It is being imported and used in the module app.py

1. As a context manager (e.g. in a pytest fixture):

import pytest


@pytest.fixture
def patch_decorator():

    def decorator_replacement(func):
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)
        return wrapper

    with PatchDecorator(
        source_module_name="utils",
        decorator_name="decorator",
        destination_module_name="app",
        replacement_decorator_obj=decorator_replacement,
    ) as patched_destination_module:
        yield patched_destination_module

2. Without context manager (not recommended):

try:
    patch_decorator = PatchDecorator(
        source_module_name="utils",
        decorator_name="decorator",
        destination_module_name="app",
        replacement_decorator_obj=decorator_replacement,
    ).setup()
finally:
    patch_decorator.cleanup()

This would also work the same for a decorator defined in a module inside a package:

import pytest


@pytest.fixture
def patch_decorator():

    def decorator_replacement(func):
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)
        return wrapper

    with PatchDecorator(
        source_module_name="package.utils",  # here
        decorator_name="decorator",
        destination_module_name="app",
        replacement_decorator_obj=decorator_replacement,
    ) as patched_destination_module:
        yield patched_destination_module

N.B: The recipe would not work if the decorator is defined and used in the same module.