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:

  1. Logging: Recording method calls, arguments, and return values for debugging or auditing
  2. Performance Monitoring: Measuring execution time of methods
  3. Access Control: Adding authentication or authorization checks
  4. Input Validation: Validating method arguments
  5. Caching: Caching method results
  6. Transaction Management: Managing database transactions
  7. 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:

  1. Regular Instance Methods:

    def method(self, arg):
    pass

  2. Class Methods:

    @classmethod
def method(cls, arg):
    pass

  3. Static Methods:

    @staticmethod
def method(arg):
    pass

  4. Properties:

    @property
def prop(self):
    pass

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

  5. 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
)

# 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:

  1. Ensure type safety
  2. Make the code more maintainable
  3. Provide better IDE support
  4. 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:

  1. Preserves special attributes from the original function
  2. Stores the wrapper function and its arguments
  3. Determines if we’re wrapping a free function or a method
  4. 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

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

Example Usage with EmailStr Class

class EmailStr:
    """A class for handling email addresses with validation."""
    
    def __init__(self, address: str):
        """Initialize with an email address."""
        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) -> Tuple[str, str]:
        """Validate an email address and return its parts."""
        import re
        email_regex = re.compile(
            r"^(?P<local_part>[a-zA-Z0-9._%+-]+)"
            r"@"
            r"(?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")

    @property
    def domain_extension(self) -> str:
        """Get the top-level domain."""
        return self.domain_part.split(".")[-1]

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

    @classmethod
    def from_parts(cls, local_part: str, domain_part: str) -> "EmailStr":
        """Create an EmailStr from parts."""
        return cls(f"{local_part}@{domain_part}")

# Decorate all methods with logging
for k, v in EmailStr.__dict__.items():
    if k not in object.__dict__ and hasattr(v, "__get__"):
        setattr(EmailStr, k, wrapping(v, wrapper_fn=log_fn(writer.write)))

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

  1. Python Descriptor Protocol Documentation
  2. Python Data Model Reference
  3. Stack Overflow: Wrap all methods of python superclass
  4. Stack Overflow: Completely wrap an object in Python
  5. Python Type Hints Documentation

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