Python: Advanced Decorators#

In a previous post, I shared a primer on how to approach the thinking of decorators and when we can apply them in our code. To summarize, we primarily see 3 entry points where decorators can dynamically effect our code:

entry point

attainable by

function/class definition

decorator only

prior to function/class call

decorator + wrapper

after function/class call

decorator + wrapper

This week we’ll revisit how we can take further advantage of these entry points in order to apply some more complex dynamic behavior.

Decorating a class#

Why would we ever want to decorate a class? Well the primary reason is, like with any decorator- we want to uniformly and dynamically update the objects state or behavior. Or perhaps we want to register all instances of the class or the class itself to some other object.

Now, if you’re thinking that these sound like patterns you can achieve with meta classes or the __init_subclass__ dunder method- then you’re spot on! Decorators give us a slightly different way to manipulate classes at their definition or instantiation in the same fashion as functions. But when would you want to use a decorator over __init_subclass__? The answer lies in whether or not you want to necessitate inheritance. To use either meta classes or __init_subclass__ you’ll need to rely on inheritance, necessitating to definition of another parent or meta class.

Note

PEP 487 – Simpler customisation of class creation introduced __init_subclass__ as a means to simplify a common pattern achieved via metaclasses. The idea was to make this pattern attainable without forcing the user to need to understand the complexity metaclasses introduce.

Another concern is that if you intend to return a wrapper to a class then you end up adding a layer of abstraction between your class and its instantiation. By returning a wrapper function from a decorated class, the user will no longer have direct access to that class or any of it’s methods unless you return a wrapper class- and if that’s the case I would probably lean towards the more traditional object-oriented approach of inheritance to solve the problem at hand.

This question ultimately boils down to what type of behavior I want to dynamically add to my class, and how easy is the implementation to follow. Above I pointed out my opinion that can be simplified to the following flow chart

  1. Modify the class at definition?

    • __init_subclass__ if this behavior should propagate to all subclasses?

    • class decorator to avoid inheritance ‘side effects’

  2. Modify the class at instantiation?

    • Inheritance > decorator that returns class wrapped

We’ve discussed some of the motivating theory of decorating a class, but let’s take a look at something a little more tangible and write one ourselves!

Writing a class decorator#

A class decorator can be written the exact same way as a function decorator- just know that the object that is passed into your decorating function will be the decorated class instead of a function. We can manipulate the passed class in any manner that Python allows us to- including the dynamic addition of methods!

def add_foo_method(cls):
    def _foo(self, base=0):
        print(f'called from {self.__class__.__name__}')
        return base + 5

    cls.foo = _foo
    return cls

@add_foo_method
class A:
    pass

@add_foo_method
class B:
    pass
a = A()
a.foo(10)
called from A
15
b = B()
b.foo(30)
called from B
35

Here you can see I defined a decorator that takes a class as its argument (instead of a function). I then rely on the decorator to dynamically add a method to that class on class definition. Note that even though I have defined a function within my decorator, it is not a wrapper since it does not wrap the passed class or function.

So this is a function that decorates a class in order to alter that class dynamically. If you’ve ever worked with dataclasses, this is exactly how @dataclass works under the hood (though much more than a single method is added).

But what about the opposite pattern? Can we use a class as a decorator?

Class as a decorator#

Given that decorators work by simply calling a decorating function, why can’t we use a class as a decorator? If the only thing that occurs with the decorating object is calling it, then we should be able to add some logic

class Registry:
    all_funcs = []

    def __init__(self, func):
        self.func = func
        self.all_funcs.append(func)

    @classmethod
    def execute_registered(cls):
        for f in cls.all_funcs:
            print(f'{f.__name__}() = {f()}')

@Registry
def f1():
    return 1

@Registry
def f2():
    return 2

@Registry
def f3():
    return 3

print(Registry.all_funcs)
Registry.execute_registered()
[<function f1 at 0x7f69dac6b7f0>, <function f2 at 0x7f69dac6ba30>, <function f3 at 0x7f69dac6bac0>]
f1() = 1
f2() = 2
f3() = 3

In the above you can see me use a class as a decorator to track a group of functions. I can then use a method of that class to execute all of the functions! I should mention that we don’t see classes used as decorators outright very often. This relies on the __init__ method to accept a function as its argument and store them into all_funcs class attribute.

But there is something quite wrong with the above snippet. Let’s take a look at the returned decorated functions-

print(f1, f2, f3, sep='\n')
<__main__.Registry object at 0x7f69dacf7130>
<__main__.Registry object at 0x7f69dacf7d30>
<__main__.Registry object at 0x7f69dacf4b20>

They’ve all become instances of the Registry class! Considering all we wanted to do was perform a registration from each of these functions to some shared state, we’ve introduced quite a large side effect. If we want to use the functions as written, we now need to reach inside of the returned instance (e.g. f1.func()) to access f1 as written.

This is because we did not pass the function through. In fact, passing any object through by means of just an __init__ method is impossible! This is because __init__ must return None. There are a few ways to implement this passthrough behavior to make sure that each decorated function is itself (instead of an instance of the decorated class). One way simply replace the __init__ definition with an implementation of __new__. This will enable us to work with the class attributes for registration and return the decorated function instead of returning an instance of the class.

class Registry:
    all_funcs = []

    def __new__(cls, func):
        cls.all_funcs.append(func)

        # pass the function through instead
        #   of returning an instance of this class
        return func

    @classmethod
    def execute_registered(cls):
        for f in cls.all_funcs:
            print(f'{f.__name__}() = {f()}')

@Registry
def f1():
    return 1

@Registry
def f2():
    return 2

@Registry
def f3():
    return 3

print(f1, f2, f3, sep='\n')
Registry.execute_registered()
<function f1 at 0x7f69dac6b520>
<function f2 at 0x7f69dac6bc70>
<function f3 at 0x7f69dac6bd00>
f1() = 1
f2() = 2
f3() = 3

Now that the functions are being passed through our decorator, they are returned as expected. Though seeing this type of pattern is not very common, it highlights the flexibility that decorators share: the decorating object just needs to be callable - not necessarily have a __call__ method.

I also want to point out that a more commonly observed pattern when dealing with classes as decorators, is to not use the class itself as the decorator, but one of its methods (e.g. @Registry.register) You see this pattern in the popular 3rd party library flask (amongst others). However before we can journey into that implementation I want to touch on one more usage of decorators: higher order decorators.

Higher Order Decorators#

These decorators are ones that accept arguments prior to decorating a class/function. You can think of a high order decorator as a function that returns a decorator. Let’s try to implement one below with a more advanced registration pattern:

from collections import defaultdict

REGISTRY = defaultdict(list)

def register(group):
    def decorator(f):
        REGISTRY[group].append(f)
        return f
    return decorator

@register('group_1')
def f1():
    pass

@register('group_1')
def f2():
    pass

@register('group_2')
def f3():
    pass

@register('group_3')
def f4():
    pass

REGISTRY
defaultdict(list,
            {'group_1': [<function __main__.f1()>, <function __main__.f2()>],
             'group_2': [<function __main__.f3()>],
             'group_3': [<function __main__.f4()>]})

Regardless of which In the above snippet, I define a global variable REGISTRY then use a high order decorator (a function that returns a decorator) to register the decorated functions to my global variable. I use a higher order decorator here to specifically categorize the decorated functions into the global REGISTRY.

To better understand this pattern, it is very important to note that @decorator() is not parsed from left to right. Instead, decorator() is first executed, and then the @ latches to whatever is returned, which in this case is a proper decorator. This is why I refer to this pattern as a decorator factory function.

I also want to note that a higher order decorator can be defined as a class as well- relying on __init__ to handle any factory stateful set up, and __call__ to be the actual decorator. This enables us to not work with a true global variable, and instead work with a class attribute instead:

from collections import defaultdict

class register:
    registry = defaultdict(list)

    def __init__(self, group):
        self.group = group

    def __call__(self, func):
        self.registry[self.group].append(func)
        return func

# Note that code from here below does not change
#     from the above example!
@register('group_1')
def f1():
    pass

@register('group_1')
def f2():
    pass

@register('group_2')
def f3():
    pass

@register('group_3')
def f4():
    pass

register.registry
defaultdict(list,
            {'group_1': [<function __main__.f1()>, <function __main__.f2()>],
             'group_2': [<function __main__.f3()>],
             'group_3': [<function __main__.f4()>]})

Ambiguously parameterized decorators#

I couldn’t wrap this post on decorators without discussing a divisive pattern of decorators.

As made convenient by decorator.decorator, and in the source code for @dataclass there is a way to implement a decorator such that either take parameterized arguments, or no arguments and ‘just work’. I use this term loosely because I personally don’t like the ambiguity introduced by this design pattern. Nonetheless, in the spirit of Python: we’re all consenting adults.

def maybe_dec_factory(f=None, /, *, kwarg1=None, kwarg2=None):
    def dec(f):
        print(f'decorating {f.__name__} with {kwarg1 = }, {kwarg2 = }')

        def wrapper(*args, **kwargs):
            print(f'before {f}')
            result = f(*args, **kwargs)
            print(f'after {f}')
            return result
        return wrapper

    if f is None:
        return dec
    return dec(f)

@maybe_dec_factory
def f1():
    print('executing f1')

@maybe_dec_factory(kwarg1='hello', kwarg2='world')
def f2():
    print('executing f2')
decorating f1 with kwarg1 = None, kwarg2 = None
decorating f2 with kwarg1 = 'hello', kwarg2 = 'world'
f1()
before <function f1 at 0x7f69d835c280>
executing f1
after <function f1 at 0x7f69d835c280>
f2()
before <function f2 at 0x7f69d835c4c0>
executing f2
after <function f2 at 0x7f69d835c4c0>

In the above example, you’ll notice that I can use this decorator both with and without calling it.

  • @maybe_dec_factory

  • @maybe_dec_factory(kwarg1='hello', kwarg2='world')

This enables me to pass optional parameters to the decorator, and avoid the following syntax:

  • @maybe_dec_factory()

I dislike this style because it results in decorators that are ambiguous in their call signature. Additionally, issues can arise if a user decides to pass some positional argument that is not the to be decorated function itself. Here instead of using keyword arguments, I attempt to pass a positional argument, thus taking the place of f in my maybe_dec_factory function. My code is currently set up under the assumption that f is a function, whereas in this case it will be the string 'hello' and will thus lead to an error.

@maybe_dec_factory('hello')
def f3():
    print('executing f3')
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Input In [12], in <cell line: 1>()
----> 1 @maybe_dec_factory('hello')
      2 def f3():
      3     print('executing f3')

Input In [9], in maybe_dec_factory(f, kwarg1, kwarg2)
     12 if f is None:
     13     return dec
---> 14 return dec(f)

Input In [9], in maybe_dec_factory.<locals>.dec(f)
      2 def dec(f):
----> 3     print(f'decorating {f.__name__} with {kwarg1 = }, {kwarg2 = }')
      5     def wrapper(*args, **kwargs):
      6         print(f'before {f}')

AttributeError: 'str' object has no attribute '__name__'
class T:
    def __init__(self, func):
        pass

    @classmethod
    def decorate(cls, func):
        return cls(func)

def decorate(func):
    return T(func)

# 1) decorator = function
@decorate
def f1():
    return 1
print(f1)

# 2) decorator = class
@T
def f1():
    return 1
print(f1)

# 3) decorator = class method
@T.decorate
def f1():
    return 1

print(f1)
<__main__.T object at 0x7f69c97083d0>
<__main__.T object at 0x7f69c97098d0>
<__main__.T object at 0x7f69c97083d0>

Summary#

And that takes us to the end of this post! I hope you enjoyed it and are ready to use some of these concepts in your own work. Remember to use decorators responsibly!

  • When decorating a class to determine if you actually want a decorator or some type of inheritance/meta class solution.

  • Classes themselves can also be used as decorators, but when doing so make sure you’re tapping into the correct hooks in Python’s object model.

  • Both functions and classes can be used to create parameterized decorators introducing some additional variables that can be accessed by the returned decorator.

    • This is useful to create a common decorator-based entry point that can dynamically change the behavior of the decorated object.

    • These types of decorators can be written to either take arguments or not- it’s up to you to to decide whether increased ambiguity is worth the slightly less verbose syntax

Talk to you all next week!