23.6 Class-Based Decorators
While function-based decorators are common, class-based decorators offer a more structured and powerful approach, particularly for stateful decorators or those requiring complex configuration. A class becomes a decorator by making its instances callable, typically by implementing the __call__ method. This method is invoked whenever the decorated function is called, allowing the class to intercept, modify, or replace the call.
Implementing a Basic Class-Based Decorator
The fundamental mechanism is to have the class’s __init__ method accept and store the function to be decorated. The __call__ method then wraps the original function, executing code before and after its invocation.
class TraceCalls:
"""A decorator class that logs every call to the decorated function."""
def __init__(self, func):
# This is called at decoration time (@TraceCalls)
self.func = func # Store the original function
# Optional: copy metadata from func to the instance
# functools.wraps does this for functions, we do it manually for classes
self.__name__ = func.__name__
self.__doc__ = func.__doc__
def __call__(self, *args, **kwargs):
# This is called every time the decorated function is invoked
print(f"Calling {self.func.__name__} with args: {args}, kwargs: {kwargs}")
result = self.func(*args, **kwargs) # Proceed with the original call
print(f"{self.func.__name__} returned: {result}")
return result
@TraceCalls
def add(a, b):
"""Adds two numbers."""
return a + b
result = add(3, b=5)
# Output:
# Calling add with args: (3,), kwargs: {'b': 5}
# add returned: 8
print(result) # Output: 8
Maintaining State Between Calls
A significant advantage of class-based decorators is their innate ability to maintain state across multiple calls to the decorated function. Instance variables are ideal for this purpose, as they persist for the lifetime of the decorator instance.
class CallCounter:
"""A decorator that counts the number of times a function is called."""
def __init__(self, func):
self.func = func
self.count = 0 # State is initialized once, at decoration
def __call__(self, *args, **kwargs):
self.count += 1 # State is persisted between calls
print(f"{self.func.__name__} has been called {self.count} times.")
return self.func(*args, **kwargs)
@CallCounter
def greet(name):
return f"Hello, {name}!"
greet("Alice")
greet("Bob")
# Output:
# greet has been called 1 times.
# greet has been called 2 times.
Parametrizing Class-Based Decorators
To create a parametrized decorator (e.g., @Delay(seconds=2)), the class’s __init__ method must accept parameters. The decoration syntax now creates an instance of the class. The function is passed later, not to __init__, but to the __call__ method of that instance. This requires a two-step process.
class Delay:
"""A parametrized decorator to delay function execution."""
def __init__(self, seconds=1):
# This is called when the decorator is parametrized: @Delay(seconds=2)
self.seconds = seconds # Store the parameter
def __call__(self, func):
# This is called with the function to be decorated
# It must return a callable that replaces the original function
def wrapper(*args, **kwargs):
import time
print(f"Waiting for {self.seconds} second(s)...")
time.sleep(self.seconds)
return func(*args, **kwargs)
# Copy metadata to the wrapper function
wrapper.__name__ = func.__name__
wrapper.__doc__ = func.__doc__
return wrapper # Return the new wrapper function
@Delay(seconds=2)
def say_hello():
print("Hello, world!")
say_hello() # Will wait 2 seconds before printing
Best Practices and Common Pitfalls
Preserving Metadata: Unlike functools.wraps, class-based decorators don’t automatically handle metadata (like __name__ or __doc__). You must manually copy these attributes from the original function to the callable object you return (either the instance itself or an internal wrapper function), as shown in the examples above. Failing to do so breaks introspection tools and can confuse debugging.
The __call__ vs. __init__ Pitfall: A common mistake is confusing what __init__ and __call__ receive in a parametrized decorator. Remember: __init__ gets the decorator’s parameters. __call__ (of the instance created by __init__) gets the function being decorated. The instance must be callable and return a callable that replaces the original function.
Inheritance for Composition: Class-based decorators can elegantly use inheritance to compose functionality.
class Logger:
def __init__(self, func):
self.func = func
def __call__(self, *args, **kwargs):
print(f"Log: Entering {self.func.__name__}")
return self.func(*args, **kwargs)
class Timer(Logger):
"""A decorator that combines logging and timing by inheriting from Logger."""
def __call__(self, *args, **kwargs):
import time
start_time = time.perf_counter()
# Use super() to call the parent's __call__ method (which includes logging)
result = super().__call__(*args, **kwargs)
end_time = time.perf_counter()
print(f"Execution time: {end_time - start_time:.4f} seconds")
return result
@Timer
def calculate_sum(n):
return sum(range(n))
calculate_sum(1000000)