How to Write Custom Python Decorators

In Python, a decorator allows a user to add useful functionalities to existing object. To use a decorator, we use the @symbol followed by the name of a decorator.

@mydecorator
def myfunction():
    pass

When calling myfunction(), the decorator mydecoratoris called first before executing myfunction.

Decorators can be stacked. In the example below, when we try to call myfunction(), mydecorator1 is called first, followed by mydecorator2 and then myfunction is called last.

@mydecorator2
@mydecorator1
def myfunction():
    pass

Decorators can also accept arguments similar to a function.

@mydecorator(5)
def myfunction():
    pass

Decorators can also be used with classes.

@mydecorator
class MyClass:
    pass

Writing custom decorators

But how do we write our own custom decorators? Commonly, decorators are just functions. There are other "creative" ways to write a decorator but we will simply stick to functions. Here are 3 useful examples.

Decorator without arguments

Let's say we want to create a decorator for caching results from functions. First, we will use the following pattern.

def mydecorator(f):
    def inner_function(*args, **kwargs):
        return f(*args, **kwargs)
    return inner_function

In the above example, inner_function is nested within mydecorator. This is how decorators are written. mydecorator accepts the function object f as an argument which points to the function where we are applying the decorator to. The purpose of inner_function is to accept the arguments intended for function f. On line 3 is the actual function call, return f(*args, **kwargs). Take note that mydecorator only return the inner_function object and does not call it.

One example application is for creation of a "caching" decorator. A caching decorator is useful for slow functions where we cache the results of the function for specific argument it accepted. This is useful if the result of the function will never change or not affected by other factors aside from the arguments themselves. Below is a naive implementation.

# decorators.py
import json
import logging
from collections import OrderedDict
from copy import copy

logger = logging.getLogger(__name__)
CACHE = {}


def sort(kwargs):
    sorted_dict = OrderedDict()
    for key, value in kwargs.items():
        if isinstance(value, dict):
            sorted_dict[key] = sort(value)
        else:
            sorted_dict[key] = value
    return sorted_dict


def cached(f):
    def _cached(*args, **kwargs):
        cache_key = json.dumps([id(f), args, sort(kwargs)], separators=(',', ':'))
        if cache_key in CACHE:
            logger.info("returning cached result")
            return CACHE[cache_key]
        result = f(*args, **kwargs)
        logger.info("caching result")
        CACHE[cache_key] = copy(result)
        return result
    return _cached

In this example, we use the function ID, its positional arguments args and its keyword arguments kwargs as key to our cached data. We converted it into a JSON string because a Python dictionary may not be able to hash some objects (e.g. a dict is not hashable). We sorted the keyword arguments recursively to make sure unsorted dictionaries will result in the same value.

To use the cached decorator, just add @cached to a function declaration.

import logging
import sys

from decorators import cached

logging.basicConfig(stream=sys.stdout, level=logging.INFO)


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


if __name__ == '__main__':
    print(add(5, 10))
    print(add(5, 10))
    print(add(5, b=10)

We added logging to check what is happening internally. Notice that on the 3rd function call, it did not use the cached data. What happened here is that b was included in kwargs and not in args.

INFO:decorators:caching result
15
INFO:decorators:returning cached result
15
INFO:decorators:caching result
15

Decorator with arguments

Decorators may accept arguments like normal functions do. In writing this type of decorator, we simply add another nested function to our original pattern.

def mydecorator(*args, *kwargs):
    def inner_function(f):
        def innermost_function(*args, **kwargs):
            return f(*args, **kwargs)
        return innermost_function
    return inner_function

The difference now is that the outermost function mydecorator will now accept custom arguments for the decorator, inner_function will accept the function object and the innermost_function with accept the arguments intended for the decorated function.

Improving on our caching decorator earlier, we want to limit the maximum number of cached data as an example. What we will be implementing is that old keys will be removed when new data are added and it exceeds the limit.

# imports

CACHE = {}
KEYS = []

# sort function

def cached(limit=100):
    def _cached(f):
        def __cached(*args, **kwargs):
            cache_key = json.dumps([id(f), args, sort(kwargs)], separators=(',', ':'))
            if cache_key in CACHE:
                logger.info("returning cached result")
                return CACHE[cache_key]
            result = f(*args, **kwargs)
            logger.info("caching result")
            CACHE[cache_key] = copy(result)
            KEYS.append(cache_key)
            if len(KEYS) > limit:
                key = KEYS.pop(0)
                logger.info("removed key %s", key)
            return result
        return __cached
    return _cached

Let us test our decorator to a limit of 3 items only using the code below.

@cached(3)
def add(a, b):
    return a + b


if __name__ == '__main__':
    print(add(1, 2))
    print(add(3, 4))
    print(add(5, 6))
    print(add(7, 8))

After executing the above code, we will notice that the first item was removed on the 4th add() call.

3
INFO:decorators:caching result
7
INFO:decorators:caching result
11
INFO:decorators:caching result
INFO:decorators:removed key [4341464816,[1,2],{}]
15

Decorator for classes

Writing decorator for classes is the same as writing decorator for functions. In this example, we will emulate the extends functionality found in other programming languages. extends is simply an inheritance pattern.

def extends(base):
    def _extends(original_class):
        class ExtendedClass(original_class, base):
            pass
        return ExtendedClass
    return _extends

In the above example, we simply created a new class from 2 parameters, base and original_class. Because we do not need any additional arguments for a class, we only have a two-level nested function.

As an example, let's extend the class Bar with Foo:

from decorators import extends


class Foo:
    def foo(self):
        return "foo"


@extends(Foo)
class Bar:
    def bar(self):
        return "bar"


if __name__ == '__main__':
    a = Bar()
    print(a.foo())
    print(a.bar())

The above example will output:

foo
bar

Conclusion

Writing custom decorators is simple and fun to do. We also learned that we can nest functions and dynamic inheritance is possible.

Read more articles like this in the future by buying me a coffee!

Buy me a coffee