Deep Dive into Python Class Decorator

Photo by Vicky Ng on Unsplash

In my previous blog, we looked at what a python decorator is and how to use it. In this blog, we are going to look at how to decorating a class instead of simply a function. We also look at a special kind of it — a class decorator.

Two ways of decorating Classes

To decorate a class, you can use two ways. One is to decorate the individual method inside the class, e.g. using @classmethod, @staticmethod, and @property.

The other way, is to decorate the whole class itself, e.g. the @dataclass introduced in python 3.7.

from dataclasses import dataclass
@dataclass
class A:
    a: str

The first way is useful when you only need to decorate some methods, but the second way is useful when you want to apply some constraints on the whole class.

Decorating a class

A common use of class decorators is to replace some usage of metaclasses.

When function decorators were originally debated for inclusion in Python 2.4, class decorators were seen as obscure and unnecessary thanks to metaclasses. After several years’ experience with the Python 2.4.x series of releases and an increasing familiarity with function decorators and their uses, the BDFL and the community re-evaluated class decorators and recommended their inclusion in Python 3.0

The motivating use-case was to make certain constructs more easily expressed and less reliant on implementation details of the CPython interpreter. While it is possible to express class decorator-like functionality using metaclasses, the results are generally unpleasant and the implementation highly fragile . In addition, metaclasses are inherited, whereas class decorators are not, making metaclasses unsuitable for some, single class-specific uses of class decorators.

— rational PEP

A class decorator is another decorator that receives a class instance as a parameter instead of a function and returns a modified version of the class. Let’s look at an example:

import functools
def decor(cls):
    @functools.wraps(cls)
    def wrapper_decor(*args, **kwargs):
        print('Echoing func %s\n' % cls.__name__)
        print('Echoing args %\n' % args)
        print('Echoing kwargs %s\n' % kwargs)
        return wrapper_decor = cls(*args, **kwargs)
    return wrapper_decor
@decor
class A:
    pass

Construct a “class” decorator

A common class decorator is a decorator to decorate a class, but you can also write a class decorator. A class decorator is especially useful when you want to have a stateful decorator as you can save the state in the class variables.

To write a “class” decorator, you will need to implement .__init__() that takes the decorated func as an argument and a .__call__() method in order for the decorator class to be callable by the decorated class.

import functools
class Count:
    def __init__(self, func):
        functools.update_wrapper(self, func)
        self.func = func
        self.count = 0
def __call__(self, *args, **kwargs): 
        self.count += 1
        print(f"Number of calls: {self.count}")
        return self.func(*args, **kwargs)
@Count
def just_example():
    print("See how many calls")
just_example()
just_example()
just_example()
====================>
See how many calls
1
See how many calls
2
See how many calls
3

Note the functools.update_wrapper() function , is essentially a full version of @functools.wraps that we often use with regular decorator:

functools.update_wrapper(wrapper, wrapped, assigned=WRAPPER_ASSIGNMENTS, updated=WRAPPER_UPDATES)

Update a wrapper function to look like the wrapped function. The optional arguments are tuples to specify which attributes of the original function are assigned directly to the matching attributes on the wrapper function and which attributes of the wrapper function are updated with the corresponding attributes from the original function. The default values for these arguments are the module level constants WRAPPER_ASSIGNMENTS (which assigns to the wrapper function’s __module__, __name__, __qualname__, __annotations__ and __doc__, the documentation string) and WRAPPER_UPDATES (which updates the wrapper function’s __dict__, i.e. the instance dictionary).

The main intended use for this function is in decorator functions which wrap the decorated function and return the wrapper. If the wrapper function is not updated, the metadata of the returned function will reflect the wrapper definition rather than the original function definition, which is typically less than helpful.

That’s so much of it!

Happy Reading!