Class as Decorator in Python

While function-based decorators are widely used, class-based decorators can offer more flexibility and better organization, especially when the decorator needs to maintain state or requires multiple methods to function properly. A Python class decorator is simply a class that implements the __call__ method, allowing an instance of the class to be used as a decorator.

This article explores the concept of using classes as decorators in Python, illustrating their purpose, implementation, and benefits through practical examples.

Implementing Class Decorators

Example: Let’s consider a simple example where we use a class decorator to log the execution time of a function:

import time

class TimerDecorator:
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        start_time = time.time()
        result = self.func(*args, **kwargs)
        end_time = time.time()
        print(f"Function {self.func.__name__} executed in {end_time - start_time} seconds")
        return result

@TimerDecorator
def example_function(n):
    total = 0
    for i in range(n):
        total += i
    return total

# Usage
print(example_function(1000000))

Function example_function executed in 0.06820821762084961 seconds
499999500000

Class Decorator with *args and **kwargs

In this example, the class decorator demonstrates how you can modify the return value of a function, adding additional logic to its execution.

Example:

1. LoggerDecorator Class:
   • The __init__ method takes the function to be decorated and stores it.
   • The __call__ method logs the arguments and keyword arguments, calls the original function with these arguments, and returns the result.
2. greet Function:
   • This is a simple function that takes a name and an optional greeting message.
3. Using the Decorator:
   • When greet is called, the LoggerDecorator logs the arguments before executing the function.

class LoggerDecorator:
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        print(f"Arguments: {args}, Keyword Arguments: {kwargs}")
        result = self.func(*args, **kwargs)
        return result

@LoggerDecorator
def greet(name, greeting="Hello"):
    return f"{greeting}, {name}!"

# Usage
print(greet("Alice"))
print(greet("Bob", greeting="Hi"))

Arguments: ('Alice',), Keyword Arguments: {}
Hello, Alice!
Arguments: ('Bob',), Keyword Arguments: {'greeting': 'Hi'}
Hi, Bob!

Class Decorator with return Statement

In this example, the class decorator demonstrates how you can modify the return value of a function, adding additional logic to its execution.

Example:

1. DoubleReturnDecorator Class:
   • The __init__ method takes the function to be decorated and stores it.
   • The __call__ method calls the original function with the provided arguments, then doubles the result and returns it.
2. add Function:
   • This is a simple function that adds two numbers.
3. Using the Decorator:
   • When add is called, the DoubleReturnDecorator modifies the return value by doubling it before returning it.

class DoubleReturnDecorator:
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        result = self.func(*args, **kwargs)
        return result * 2

@DoubleReturnDecorator
def add(a, b):
    return a + b

# Usage
print(add(3, 5))  
print(add(10, 20)) 

16
60

Checking for an Invalid Parameter with a Class Decorator

In this example, the class decorator demonstrates how you can check for specific parameter values and handle errors appropriately by raising exceptions when conditions are not met.

Example

1. ErrorCheckDecorator Class:
   • The __init__ method takes the function to be decorated and stores it.
   • The __call__ method checks if the error keyword argument is set to True. If it is, it raises a ValueError. Otherwise, it calls the original function with the provided arguments.
2. process_data Function:
   • This is a simple function that takes some data and an optional error flag.
3. Using the Decorator:
   • When process_data is called with error=True, the decorator raises a ValueError. Otherwise, it processes the data normally.

class ErrorCheckDecorator:
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        if kwargs.get('error', False):
            raise ValueError("Invalid parameter 'error' set to True")
        return self.func(*args, **kwargs)

@ErrorCheckDecorator
def process_data(data, error=False):
    return f"Processing data: {data}"

# Usage
try:
    print(process_data("sample_data"))
    print(process_data("sample_data", error=True))  # This will raise ValueError
except ValueError as e:
    print(e)

Processing data: sample_data
Invalid parameter 'error' set to True