All Decorators - Systematically Decorating Python Class Methods

A deep dive into systematically decorating all methods of a Python class, exploring challenges and solutions using the Descriptor protocol.

Introduction#

Decorators are a powerful feature in Python, allowing you to modify the behavior of functions or methods. They provide a clean and elegant way to wrap functions with additional functionality without modifying their original code. However, when it comes to decorating all methods of a class, several challenges arise that make this seemingly simple task quite complex.

Understanding the Problem Space#

Before diving into the solution, let's understand why decorating all methods of a class is challenging and why existing solutions might not be sufficient.

The Need for Method Decoration#

There are many scenarios where you might want to decorate all methods of a class:

Logging: Recording method calls, arguments, and return values for debugging or auditing
Performance Monitoring: Measuring execution time of methods
Access Control: Adding authentication or authorization checks
Input Validation: Validating method arguments
Caching: Caching method results
Transaction Management: Managing database transactions
Error Handling: Adding consistent error handling across methods

Current Approaches and Their Limitations#

Manual Decoration#

The most straightforward approach is manually decorating each method:

class MyClass:
    @decorator
    def method1(self):
        pass
    
    @decorator
    def method2(self):
        pass

Limitations:

Tedious and error-prone
Easy to forget to decorate new methods
Clutters the code
Hard to maintain consistency

Metaclass Approach#

Another common approach is using metaclasses:

class DecoratorMetaclass(type):
    def __new__(cls, name, bases, attrs):
        for attr_name, attr_value in attrs.items():
            if callable(attr_value):
                attrs[attr_name] = decorator(attr_value)
        return super().__new__(cls, name, bases, attrs)

class MyClass(metaclass=DecoratorMetaclass):
    def method1(self):
        pass

Limitations:

Doesn't handle all method types correctly
Can interfere with method resolution order
Makes class inheritance more complex
Not intuitive for Python developers

The Challenges#

Method Type Diversity#

Python has various types of methods, each with its own characteristics:

Regular Instance Methods:
```
def method(self, arg):
    pass
```

Class Methods:

@classmethod
def method(cls, arg):
    pass

Static Methods:

@staticmethod
def method(arg):
    pass

Properties:

@property
def prop(self):
    pass

@prop.setter
def prop(self, value):
    pass

Abstract Methods:

@abstractmethod
def method(self):
    pass

Each of these method types implements the descriptor protocol differently, making it challenging to create a universal wrapper.

Descriptor Protocol Complexity#

The descriptor protocol is what makes properties, methods, and other decorators work in Python. A descriptor is an object attribute with "binding behavior", where attribute access is overridden by methods in the descriptor protocol. The protocol includes:

__get__(self, obj, type=None)
__set__(self, obj, value)
__delete__(self, obj)

Understanding how each method type implements these methods is crucial for creating a universal wrapper.

The Solution: Using the Descriptor Protocol#

Our solution leverages the descriptor protocol to create a universal wrapper that can handle any method type. Let's break down the implementation step by step.

Type Definitions#

First, we define our types using Python's typing system:

from typing import (
    Any, Callable, Dict, Generic, Optional, 
    Tuple, Type, TypeVar, Union, cast
)
import json
from functools import wraps
from dataclasses import dataclass

# Contravariant type for method binding
T_contra = TypeVar("T_contra", bound=Any, contravariant=True)
# Regular type variable
T = TypeVar("T", bound=Any)
# Covariant type for return values
U_co = TypeVar("U_co", covariant=True)

# Type aliases for clarity
Func = Callable[..., U_co]
DunderGetFunc = Callable[[T_contra, Type[T_contra]], Func]
DunderGetFuncOrFunc = Union[DunderGetFunc, Func]
WrapperFunc = Callable[[Func, Tuple[Any, ...], Dict[str, Any]], Func]

These type definitions help us:

Ensure type safety
Make the code more maintainable
Provide better IDE support
Document the expected types

Constants and Utilities#

We define constants for special method names and a utility function:

class LiteralConsts:
    """Constants for special method names in Python."""
    DUNDER_GET = "__get__"
    DUNDER_DOC = "__doc__"
    DUNDER_ABSTRACT = "__isabstractmethod__"
    DUNDER_MODULE = "__module__"
    DUNDER_NAME = "__name__"
    DUNDER_QUALNAME = "__qualname__"

def pass_though(x: T, _args, _kwargs) -> "T":
    """Default wrapper that simply returns the input unchanged."""
    return x

The WrapperFn Class#

The heart of our solution is the WrapperFn class:

class WrapperFn(Generic[T_contra, U_co]):
    """
    A class that wraps a function and allows additional behavior to be injected.
    
    This class implements the descriptor protocol to handle all types of methods
    correctly, including regular methods, static methods, class methods, and
    properties.
    """

Initialization#

The __init__ method handles setting up the wrapper:

def __init__(
    self,
    fn: DunderGetFuncOrFunc,
    wrapper_fn: WrapperFunc = pass_though,
    wrapper_fn_args: Optional[Tuple[Any, ...]] = None,
    wrapper_fn_kwargs: Optional[Dict[str, Any]] = None,
):
    # Copy special attributes from the original function
    for k in (
        LiteralConsts.DUNDER_ABSTRACT,
        LiteralConsts.DUNDER_DOC,
        LiteralConsts.DUNDER_MODULE,
        LiteralConsts.DUNDER_NAME,
        LiteralConsts.DUNDER_QUALNAME,
    ):
        v = getattr(fn, k, None)
        if v is not None:
            setattr(self, k, v)

    # Store wrapper function and its arguments
    self._wrapper_fn = wrapper_fn
    self._wrapper_fn_args = wrapper_fn_args if wrapper_fn_args is not None else ()
    self._wrapper_fn_kwargs = wrapper_fn_kwargs if wrapper_fn_kwargs is not None else {}

    # Determine if this is a free function or a method
    self._is_free_fn = isinstance(fn, (type(print), type(pass_though)))
    if not self._is_free_fn:
        self._fn = cast(DunderGetFunc, getattr(fn, LiteralConsts.DUNDER_GET))
    else:
        self._fn = cast(Func, fn)

Key points about initialization:

Preserves special attributes from the original function
Stores the wrapper function and its arguments
Determines if we're wrapping a free function or a method
Handles the descriptor protocol appropriately

The Descriptor Protocol Implementation#

The __get__ method implements the descriptor protocol:

def __get__(self, inst: Optional[Any], owner: Optional[Any]):
    """
    Implements the descriptor protocol.
    
    Args:
        inst: The instance that the method is being accessed from
        owner: The class that owns the method
    
    Returns:
        The wrapper instance with bound instance and owner
    """
    self._inst = inst
    self._owner = owner
    self._is_free_fn = False
    return self

This method is called when the attribute is accessed and handles method binding.

Method Calling#

The __call__ method handles the actual function execution:

def __call__(self, *args: "Any", **kwargs: "Any") -> "U_co":
    """
    Calls the wrapped function with the provided arguments.
    
    This method handles:
    1. Free functions
    2. Bound methods
    3. Unbound methods
    4. Static methods
    5. Class methods
    6. Properties
    """
    if self._wrapper_fn_args is None:
        self._wrapper_fn_args = ()
    if not self._is_free_fn:  # not free function
        self._owner = cast(Type[Any], self._owner)
        if hasattr(self._fn, LiteralConsts.DUNDER_GET):
            # Handle descriptors (e.g., staticmethod)
            func = self._fn.__get__(self._inst, self._owner)
        else:
            # Handle bound methods
            func = self._fn(self._inst, self._owner)
        # Apply the wrapper
        func = self._wrapper_fn(
            func,
            *self._wrapper_fn_args,
            **self._wrapper_fn_kwargs,
        )
    else:
        # Handle free functions
        func = self._wrapper_fn(self._fn, *self._wrapper_fn_args, **self._wrapper_fn_kwargs)
    
    # Call the function if it's callable, otherwise return it (e.g., for properties)
    if callable(func):
        return func(*args, **kwargs)
    return func

Here is the entire WrapperFn class as a runnable snippet.

class WrapperFn(Generic[T_contra, U_co]):
    """
    A class that wraps a function and allows additional behavior to be injected.

    Attributes
    ----------
        _is_free_fn (bool): Indicates if the function is a free function.
        _fn (DunderGetFuncOrFunc): The function to be wrapped.
        _inst (Optional[Any]): The instance the function is bound to.
        _owner (Optional[Any]): The owner class of the function.
        _wrapper_fn (WrapperFunc): The wrapper function that adds additional behavior.
        _wrapper_fn_args (Optional[Tuple[Any, ...]]): Positional arguments for the wrapper function.
        _wrapper_fn_kwargs (Optional[Dict[str, Any]]): Keyword arguments for the wrapper function.
        __doc__ (Optional[str]): The docstring of the wrapped function.
        __isabstractmethod__ (Optional[str]): Indicates if the wrapped function is abstract.
    """

    def __init__(
        self,
        fn: DunderGetFuncOrFunc,
        wrapper_fn: WrapperFunc = pass_though,
        wrapper_fn_args: Optional[Tuple[Any, ...]] = None,
        wrapper_fn_kwargs: Optional[Dict[str, Any]] = None,
    ):
        """
        Initializes the WrapperFn instance.

        Args:
            fn (DunderGetFuncOrFunc): The function to be wrapped.
            wrapper_fn (WrapperFunc, optional): The wrapper function that adds additional behavior. Defaults to pass_though.
            wrapper_fn_args (Optional[Tuple[Any, ...]], optional): Positional arguments for the wrapper function. Defaults to None.
            wrapper_fn_kwargs (Optional[Dict[str, Any]], optional): Keyword arguments for the wrapper function. Defaults to None.
        """
        for k in (
            LiteralConsts.DUNDER_ABSTRACT,
            LiteralConsts.DUNDER_DOC,
            LiteralConsts.DUNDER_MODULE,
            LiteralConsts.DUNDER_NAME,
            LiteralConsts.DUNDER_QUALNAME,
        ):
            v = getattr(fn, k, None)
            if v is not None:
                setattr(self, k, v)

        self._wrapper_fn = wrapper_fn
        self._wrapper_fn_args = wrapper_fn_args if wrapper_fn_args is not None else ()
        self._wrapper_fn_kwargs = wrapper_fn_kwargs if wrapper_fn_kwargs is not None else {}

        self._is_free_fn = isinstance(fn, (type(print), type(pass_though)))
        if not self._is_free_fn:
            self._fn = cast(DunderGetFunc, getattr(fn, LiteralConsts.DUNDER_GET))
        else:
            self._fn = cast(Func, fn)

    def __get__(self, inst: Optional[Any], owner: Optional[Any]):
        """
        Gets the function bound to the instance or owner.

        Args:
            inst (Optional[Any]): The instance the function is bound to.
            owner (Optional[Any]): The owner class of the function.

        Returns
        -------
            WrapperFn: The current instance with updated binding.
        """
        self._inst = inst
        self._owner = owner
        self._is_free_fn = False
        return self

    def __call__(self, *args: "Any", **kwargs: "Any") -> "U_co":
        """
        Calls the wrapped function with the provided arguments.

        Args:
            *args (Any): Positional arguments for the wrapped function.
            **kwargs (Any): Keyword arguments for the wrapped function.

        Returns
        -------
            U_co: The result of the wrapped function call.
        """
        if self._wrapper_fn_args is None:
            self._wrapper_fn_args = ()
        if not self._is_free_fn:  # not free function
            self._owner = cast(Type[Any], self._owner)
            if hasattr(self._fn, LiteralConsts.DUNDER_GET):
                func = self._fn.__get__(self._inst, self._owner)
            else:
                func = self._fn(self._inst, self._owner)
            func = self._wrapper_fn(
                # get the bound to class or bound to instance - method
                func,
                *self._wrapper_fn_args,
                **self._wrapper_fn_kwargs,
            )  # wrap it
            # call it
        else:
            func = self._wrapper_fn(self._fn, *self._wrapper_fn_args, **self._wrapper_fn_kwargs)
        if callable(func):
            return func(*args, **kwargs)
        return func

The Wrapping Decorator#

The wrapping decorator provides a convenient interface:

def wrapping(
    func: Optional[Func] = None,
    /,
    wrapper_fn: WrapperFunc = pass_though,
    wrapper_fn_args: Optional[Tuple[Any, ...]] = None,
    wrapper_fn_kwargs: Optional[Dict[str, Any]] = None,
) -> Union[Callable[[Callable[..., Any]], WrapperFn[T_contra, U_co]], WrapperFn[T_contra, U_co]]:
    """
    A decorator that wraps a function with additional behavior.
    
    Can be used with or without arguments:
    
    @wrapping
    def func(): pass
    
    @wrapping(wrapper_fn=my_wrapper)
    def func(): pass
    """
    if func is None:
        def decorator(f: Func) -> WrapperFn[T_contra, U_co]:
            return WrapperFn(f, wrapper_fn, wrapper_fn_args, wrapper_fn_kwargs)
        return decorator
    return WrapperFn(func, wrapper_fn, wrapper_fn_args, wrapper_fn_kwargs)

Practical Example: Logging Decorator#

Let's look at a complete example of using our wrapper to create a logging decorator:

Record Class for Logging#

@dataclass(repr=False)
class Record(dict):
    """A class to store and format log records."""
    
    def __setitem__(self, key, value):
        """Ensure only allowed keys are used."""
        if key not in RECORD_ALLOWED_KEYS:
            raise TypeError(f"`Record` does not allow an entry with key = {key}")
        return super().__setitem__(key, value)

    def __repr__(self):
        """Format the record as JSON."""
        self["id"] = id(self)
        return json.dumps(self, default=lambda o: o.__repr__())

# Constants for Record
RECORD_ALLOWED_KEYS = (
    *(f"{k}_positional_argument" for k in (
        "first", "second", "third", "fourth", "fifth",
        "sixth", "seventh", "eighth", "ninth", "tenth"
    )),
    "function_name",
    "result",
    "id",
)

Logging Function#

def log_fn(sink: Optional[Callable[[Record], None]] = None):
    """
    Creates a logging decorator that records function calls.
    
    Args:
        sink: A function that consumes log records
    
    Returns:
        A decorator that wraps functions with logging
    """
    def _log_fn(fn: Callable) -> Callable:
        @wraps(fn)
        def _wrapper(*args, **kwargs) -> Callable:
            # Create a new record
            report = Record()
            report["function_name"] = fn.__name__ or fn.__qualname__ or ""

            # Log positional arguments
            for label, arg in zip(
                (f"{k}_positional_argument" for k in (
                    "first", "second", "third", "fourth",
                    "fifth", "sixth", "seventh", "eighth",
                    "ninth", "tenth"
                )),
                args
            ):
                report[label] = arg

            # Log keyword arguments
            report.update({"kwargs": kwargs or {}})

            # Call the function and log the result
            result = fn(*args, **kwargs)
            report["result"] = result

            # Send to sink if provided
            if sink:
                sink(report)

            return result
        return _wrapper
    return _log_fn

class Writeable:
    def __init__(self):
        self.buffer = []

    def write(self, message: Record):
        if not isinstance(message, Record):
            raise TypeError("`message` is not of type `Record`")
        self.buffer.append(message)

    def read(self):
        return "\n".join(repr(b) for b in self.buffer)

writer = Writeable()

Example Usage with EmailStr Class#

class EmailStr:
    def __init__(self, address: str):
        if not isinstance(address, str):
            raise TypeError("`address` must be of type `str`")
        self.address = address
        self.local_part, self.domain_part = self._validate_email(address)

    @staticmethod
    def _validate_email(address: str):
        import re

        email_regex = re.compile(
            r"^(?P<local_part>[a-zA-Z0-9._%+-]+)@(?P<domain_part>[a-zA-Z0-9.-]+\.[a-zA-Z]{2,})$"
        )
        match = email_regex.match(address)
        if not match:
            raise ValueError(f"Invalid email address: {address}")
        return match.group("local_part"), match.group("domain_part")

    @staticmethod
    def is_valid_email(address: str) -> bool:
        try:
            EmailStr._validate_email(address)
        except ValueError:
            return False
        else:
            return True

    @classmethod
    def from_parts(cls, local_part: str, domain_part: str):
        address = f"{local_part}@{domain_part}"
        return cls(address)

    @property
    def local_part(self):
        return self._local_part

    @local_part.setter
    def local_part(self, value):
        if not value:
            raise ValueError("Local part of the email cannot be empty")
        self._local_part = value

    @property
    def domain_part(self):
        return self._domain_part

    @domain_part.setter
    def domain_part(self, value):
        if not value:
            raise ValueError("Domain part of the email cannot be empty")
        self._domain_part = value

    def __repr__(self):
        return json.dumps(
            {
                "address": self.address,
                "local_part": self.local_part,
                "domain_part": self.domain_part,
            }
        )

    def get_domain_extension(self) -> str:
        """Returns the top-level domain (TLD) of the email address."""
        return self.domain_part.split(".")[-1]

    def is_gmail(self) -> bool:
        """Checks if the email address is a Gmail address."""
        return self.domain_part.lower() == "gmail.com"

    @staticmethod
    def from_str(s: str) -> "EmailStr":
        return EmailStr(s)

def test_email_str():
    # methods decorating solely with `log_fn`
    for k, v in EmailStr.__dict__.items():
        if k not in object.__dict__ and hasattr(v, "__get__"):
            if k == '__dict__': continue
            setattr(EmailStr, k, wrapping(v, wrapper_fn=log_fn(writer.write)))

    email = EmailStr("hello@bye.com")
    assert email.local_part == "hello"
    assert email.domain_part == "bye.com"

    assert EmailStr.is_valid_email("hello@bye.com")  # True
    assert not EmailStr.is_valid_email("invalid-email")  # False

    email_from_parts = EmailStr.from_parts("hello", "bye.com")
    assert email_from_parts.address == "hello@bye.com"  # hello@bye.com

    assert email_from_parts.get_domain_extension() == "com"
    assert not email_from_parts.is_gmail()

    assert EmailStr.from_str("hello@gmail.com").is_gmail()

test_email_str()
print(writer.read())

Example Output#

When using the decorated EmailStr class, you'll get log records like this:

{
    "function_name": "_validate_email",
    "first_positional_argument": "user@example.com",
    "kwargs": {},
    "result": ["user", "example.com"],
    "id": 140712834567890
}
{
    "function_name": "is_gmail",
    "kwargs": {},
    "result": false,
    "id": 140712834567891
}

Advanced Usage#

Chaining Multiple Decorators#

You can chain multiple decorators using our wrapping decorator:

def timing_wrapper(fn: Callable, *args, **kwargs) -> Callable:
    """Measure function execution time."""
    @wraps(fn)
    def wrapper(*args, **kwargs):
        import time
        start = time.time()
        result = fn(*args, **kwargs)
        end = time.time()
        print(f"{fn.__name__} took {end - start:.2f} seconds")
        return result
    return wrapper

# Apply both logging and timing
method = wrapping(
    method,
    wrapper_fn=timing_wrapper,
    wrapper_fn_args=(),
    wrapper_fn_kwargs={"log_level": "DEBUG"}
)
method = wrapping(
    method,
    wrapper_fn=log_fn(writer.write)
)

Conditional Decoration#

You can create decorators that only activate under certain conditions:

def conditional_log(condition: Callable[[], bool]):
    """Only log if condition is met."""
    def wrapper(fn: Callable, *args, **kwargs) -> Callable:
        @wraps(fn)
        def wrapped(*args, **kwargs):
            result = fn(*args, **kwargs)
            if condition():
                print(f"Called {fn.__name__} with result {result}")
            return result
        return wrapped
    return wrapper

# Only log in debug mode
import sys
debug_mode = lambda: sys.flags.debug
method = wrapping(method, wrapper_fn=conditional_log(debug_mode))

Best Practices and Considerations#

When using this decorator system, it's important to consider the performance impact. Each wrapper adds a function call overhead, so you might want to use functools.lru_cache for methods that are called frequently. Profiling your application can help ensure that the overhead remains acceptable.

Error handling is another crucial aspect. Wrappers should handle exceptions appropriately, and using context managers can aid in resource cleanup. It's essential to maintain the original function's exception contract to avoid unexpected behavior.

Documentation plays a vital role in using decorators effectively. Utilize functools.wraps to preserve function metadata, and clearly document the effects of decorators. Adding type hints to wrapper functions can also enhance code readability and maintainability.

Finally, thorough testing is necessary to ensure reliability. Test decorated methods both with and without decoration to verify that decorators preserve the intended method behavior. It's also important to test edge cases and error conditions to ensure robustness.

Conclusion#

Decorating all methods of a class using the Descriptor protocol is a game-changer for Python developers. This approach is versatile, handling everything from regular methods to static and class methods, properties, and even abstract methods. By centralizing the decoration logic, it not only keeps our codebase clean and consistent but also makes it easier to tweak behaviors without diving into each method individually.

Moreover, this method ensures that our code remains type-safe, preserving method signatures and offering robust IDE support. It's like having a safety net that allows us to experiment with different functionalities without the fear of breaking things.

The flexibility of this solution is another highlight. Whether you need to apply multiple decorators, conditionally wrap methods, or introduce custom behaviors, this approach has got you covered. It's particularly useful for adding logging, monitoring, and debugging capabilities across large applications.

In essence, this implementation is a powerful tool in the Python developer's toolkit, making complex tasks more manageable and enhancing the overall quality of the code. It's a practical solution that brings both simplicity and power to the table, making our lives as developers just a little bit easier.

References#

The complete source code for this implementation is available on Playground.

All Decorators - Systematically Decorating Python Class Methods#