Type Hinting Python Decorators, part 2

In the last blog post, we explored Python’s type hints and how they can be used with decorators, leveraging the latest versions of Python. We primarily focused on function decorators in that post. This time, our topic is specifically class decorators.

In the previous post, we also covered generic programming with the support of Python type hints. First, let’s see how to create a simple generic class data structure in Python. Below is a simple stack implementation:

class Stack:
    def __init__(self):
        self.items = []


    def push(self, item):
        self.items.append(item)


    def pop(self):
        return self.items.pop()

Now let’s see how this would look like with generic type hints:

class Stack[T]:
def __init__(self) ->None:
self.items: list[T] = []


def push(self, item: T) ->None:
self.items.append(item)


def pop(self) -> T:
returnself.items.pop()

As we see, with newer Python versions it is rather straightforward to add type hints to a data structure class that can hold elements of any arbitrary type. Now we can initialize this class as an object that supports the type we want:

int_stack = Stack[int]()
int_stack.push(1)
int_stack.push(2)
int_stack.pop() # Returns 2


str_stack = Stack[str]()
str_stack.push("a")
str_stack.push("b")
str_stack.pop() # Returns "b"

However, as one might expect, this does not pass the type checker:

int_stack = Stack[int]()
int_stack.push(1)
int_stack.push("a") # Error!

Class decorators

Class decorators offer a powerful way to enhance the functionality of classes in a reusable manner while avoiding the pitfalls of deep inheritance hierarchies. They are particularly useful for handling cross-cutting concerns, such as logging and validation, by allowing behavior to be injected without modifying the class itself. Additionally, class decorators enable dynamic runtime modifications based on conditions, such as configuration settings.

Below is a classic example of a logging class decorator that logs all the method calls of any class, as well as their return values:

from functools import wraps


def preserve_method_type(method):
    if isinstance(method, staticmethod):
        method = staticmethod(method)
    elif isinstance(method, classmethod):
        method = classmethod(method)
    return method


def make_logged_method(method_name, original_method):
    # Preserve method metadata with wraps
    @wraps(original_method)
    def logged_method(*args, **kwargs):
         print(f"Calling {method_name} with args: {args}, kwargs: {kwargs}")
         result = original_method(*args, **kwargs)
         print(f"{method_name} returned: {result}")
         return result
    return preserve_method_type(logged_method)


def add_method_logging(cls):
    for attr_name, method in cls.__dict__.items():
        # Wrap only user-defined methods
        if callable(method) and not attr_name.startswith('__'):
            # Capture each method's name and function reference 
            wrapped_method = make_logged_method(attr_name, method)
            # Replace the original method with the wrapped one
            setattr(cls, attr_name, wrapped_method)
    return cls

Example usage:

@add_method_logging
class MyClass:    
    def add(self, x: int, y: int) -> int:
        return x + y


    def multiply(self, x: int, y: int) -> int:
        return x * y

The add_method_logging decorator iterates through all attributes in the class and wraps its each user-defined method, enabling the method to log both its arguments and return value. Now it would be easy to turn the logging on and off by adding a configuration check to the decorator, for example.

The method wrapping logic is handled in a separate make_logged_method function to ensure that each method’s name and reference are properly captured. Wrapping directly within the for loop of add_method_logging would introduce a bug due to Python’s late binding: in the wrapped methods, attr_name and method would end up always referencing the final values encountered in the loop.

How, then, would we then add type-safe type hints to this? Following everything we know so far, it’s rather simple:

from collections.abc import Callable
from functools import wraps


def preserve_method_type[**P, R](method: Callable[P, R]) -> Callable[P, R]:
    if isinstance(method, staticmethod):
        method = staticmethod(method)
    elif isinstance(method, classmethod):
        method = classmethod(method)
    return method


def make_logged_method[**P, R](method_name: str, original_method: Callable[P, R]) -> Callable[P, R]:
    @wraps(original_method)
    def logged_method(*args: P.args, **kwargs: P.kwargs) -> R:
        print(f"Calling {method_name} with args: {args}, kwargs: {kwargs}")
        result = original_method(*args, **kwargs)
        print(f"{method_name} returned: {result}")
        return result
    return preserve_method_type(logged_method)


def add_method_logging[T](cls: type[T]) -> type[T]:
    for attr_name, method in cls.__dict__.items():
        # Wrap only user-defined methods
        if callable(method) and not attr_name.startswith('__'):
            # Capture each method's name and function reference 
            wrapped_method = make_logged_method(attr_name, method)
            # Replace the original method with the wrapped one
            setattr(cls, attr_name, wrapped_method)
    return cls

The only new concept introduced here is the type type hint. It allows us to specify that the argument of add_method_logging should be a class of any type T rather than an instance of type T. In this way, we can define a generic class type parameter.

Now we can take any class and wrap its methods with our logging function. The type signatures of all the methods are preserved in a generic, type-safe way.

Extending classes

Another common use case for class decorators is to add the properties of one class to another. Below is an example of a class decorator that can take any class as an argument and add its attributes and methods to any decorated class.

An advantage of this approach compared to inheritance is avoiding tight coupling and complex class hierarchies, especially in situations where the same functionalities need to be added to multiple classes. This also allows for adding features based on runtime conditions, such as logging levels.

def extend_with(extension_class):
    def extend_class(target_class):
        # Add each attribute from the extension_class to the target_class
        for attr_name, attr_value in vars(extension_class).items():
            # Skip special methods and private attributes
            if not attr_name.startswith("__"):
                setattr(target_class, attr_name, attr_value)
        return target_class
    return extend_class


# Example extension class
class Logger:
    def log(self, message):
        print(f"[LOG]: {message}")


    def error(self, message):
        print(f"[ERROR]: {message}")


# Example target class
@extend_with(Logger)
class Database:
    def connect(self):
        self.log("Connecting to the database...")
        try:
            # Simulate database connection logic
            pass
        except Exception as e: # Replace with more accurate error type
            self.error(f"Failed to connect to database, error: {e}")
            return
        self.log("Connected to database.")


# Usage
db = Database()
db.connect()

Above, extend_with is a decorator factory that takes an extension_class as an argument and returns a decorator. This returned decorator is then used to decorate a target_class, adding the attributes and methods of extension_class to it. Essentially, extend_with(SomeClass) produces a decorator that has access to SomeClass via extension_class from its enclosing scope.

When used with the Database class, extend_with is called with Logger as an argument. As a result, the produced decorator function adds the attributes and methods of Logger to Database.

How do we add type hints to this example then? The solution is not the simplest, but it looks more complicated than it is:

from typing import Callable
def extend_with[T, U](extension_class: type[U]) -> Callable[[type[T]], type[T]]:
    def extend_class(target_class: type[T]) -> type[T]:
        # Add each attribute from the extension_class to the target_class
        for attr_name, attr_value in vars(extension_class).items():
            # Skip special methods and private attributes
            if not attr_name.startswith("__"):
                setattr(target_class, attr_name, attr_value)
        return target_class
    return extend_class
class Logger:
    def log(self, message: str) -> None:
        print(f"[LOG]: {message}")
    def error(self, message: str) -> None:
        print(f"[ERROR]: {message}")
@extend_with(Logger)
class Database:
    def connect(self) -> None:
        self.log("Connecting to the database...")
        try:
            # Simulate database connection logic
            pass
        except Exception as e: # Replace with more accurate error type
            self.error(f"Failed to connect to database, error: {e}")
            return
        self.log("Connected to database.")

In this example, we define type as a generic type to represent any class. The extend_with decorator factory is structured to accept any class (extension_class) as an argument and return a callable decorator (extend_class). Here, the extension_class parameter is of type type[U], indicating it can be any class type.

The callable takes a class of type T and returns a class of the same type. The type signature of extend_class decorator reflects this behavior: it accepts target_class of type T and returns a class of the same type. This setup allows extension_class and target_class to be of different types (U and T), while ensuring that the decorator always returns a class of the same type as target_class, simply with additional functionality.

Still when the type hinting is checked with mypy, we get an error:

error: "Database" has no attribute "log"  [attr-defined]

The reason is straightforward: Python’s type system has no way of knowing that new methods have been added to Database at runtime. It only recognizes the original connect method as part of that class.

We can fix this by adding stub methods to Database:

@extend_with(Logger)
class Database:
    def log(self, message: str) -> None: ...
    def error(self, message: str) -> None: ...
    # ...

With this setup, the example passes the type checker without issues. Admittedly, writing separate type definitions for each attribute and method added from the extending class to the decorated class may seem like extra work. However, consider it this way: since we’re adding attributes and methods at runtime, there’s no way for static type hinting to account for them automatically. Defining these types explicitly bridges the gap between dynamic runtime behavior and static type checking, ensuring type safety and clarity.

Summary

In this two-part blog series, we have now covered how to use type hinting with both function and class decorators. This time, we focused on class decorators. We began by introducing generic classes and demonstrated how to use generic type parameters to type functions that accept arbitrary classes. At the same time, we presented an example decorator that adds functionality to each method of the decorated class. Finally, we explored how to add type-safe type hinting to a class decorator that extends any class with methods and attributes from another arbitrary class. In doing so, we used stub methods to align static type checking with dynamic runtime behavior.

The code examples can be found here.