Better Python Decorators with wrapt

2020-07-02 Bee likes decorated flowers.

This is a cross-post from the Scout APM blog, where I occasionally write. I also maintain the Scout Python integration, which uses wrapt for its Python decorators.

A Python decorator wraps a target function with another wrapper function. This wrapper function can add any behavior you might want. For example, it can track execution times, redefine how the wrapped function runs, or modify return values.

As a concrete example, the standard library’s functools.lru_cache() decorator wraps a target function to add caching behavior. If you call a wrapped function twice with the same arguments, the second call returns a cached value rather than re-executing the potentially costly inner function.

Scout’s Decorators

Scout’s Python integration uses decorators extensively for instrumenting third party libraries. This instrumentation allows us to measure execution time and help you find hot spots in your application.

Our instrumentation uses built-in extension mechanisms where possible, such as Django’s database instrumentation. But often libraries have no such mechanisms, so we resort to wrapping third party libraries’ functions with our own decorators. For example, we instrument jinja2’s Template.render() function with a decorator to measure template rendering time.

We value correctness of our instrumentation a lot, so that we do not affect our users’ applications. It’s important that our decorators are as transparent as possible, so they do not break assumptions held by the library or application. Therefore, we create our decorators with the wrapt library by Graham Dumpleton.

wrapt makes writing correct, transparent decorators easy. If you are writing any kind of decorators in Python, you should consider using wrapt. Let’s look at how it helps the Scout integration.

Making a Decorator with the Function Closure Pattern

Imagine we have this simple function for generating HTML, imitating Jinja2’s Template.render():

def render_html(title):
    return f"<h1>{title}</h1>"

We’d like to instrument it with a decorator to print the execution time. We can create a decorator that does this like so:

import time


def instrument(wrapped):
    def wrapper(*args, **kwargs):
        start = time.time_ns()
        result = wrapped(*args, **kwargs)
        end = time.time_ns()
        print(f"Took {end - start} nanoseconds")
        return result

    return wrapper

This example uses the “function closure pattern” that is commonly used to create decorators. The outer instrument function is called with the wrapped function as wrapped, and returns the wrapper function. THe wrapper function can access the wrapped function after instrument() has completed through the function closure.

The wrapper function:

  1. Records the start in nanoseconds with time.time_ns().
  2. Calls the wrapped function wrapped, looked up in the function closure of the outer instrument().
  3. Records the end time.
  4. Prints a message about the execution time in nanoseconds.
  5. Returns the result from calling wrapped.

You can imagine that in the Scout integration, rather than printing a message, we append details about the executed function to a log. This log is then shipped with extra details as a trace of a single web request. We aggregate this trace with those for similar requests and display the results in our GUI.

Using our Decorator

We can apply our decorator like so:

@instrument
def render_html(title):
    return f"<h1>{title}</h1>"

The outer instrument() function returns our wrapper() function, bound to the name render_html.

The @ syntax is actually a shortcut added in Python 2.4. We can expand it like so:

def render_html(title):
    return f"<h1>{title}</h1>"

render_html = instrument(render_html)

In this form, we can see that render_html is effectively defined twice: first as a plain function, second as the return value of instrument(). Because our instrument() returns our new wrapper() function, it continues to be a function. Decorators can return any Python object, although it normally doesn’t make sense to return something other than a function.

We can call our decorated render_html() with a value for title:

>>> render_html("Good Dog")
Took 1000 nanoseconds
'<h1>Good Dog</h1>'

We see the message from our wrapper printed, as well as the returned string. So our decorator works, but unfortunately it’s not very transparent. It replaces the inner render_html() function with the wrapper() function, and this causes several detectable differences.

Decorator Transparency

We can see some of differences between the unwrapped and wrapped versions in their representations (through repr()). If we execute only the function name, the interpreter shows us its representation.

Without @instrument, the function’s representation looks like this:

>>> render_html
<function __main__.render_html(title)>

With @instrument, the function representation is quite different:

>>> render_html
<function __main__.instrument.<locals>.wrapper(*args, **kwargs)>

We can see these differences in the decorated version:

These pieces of the function representation are all included from some special attributes. For example, the “qualified name” for the function is stored in its __qualname__ attribute, which is where <locals> comes from for the instrumented version:

>>> render_html.__qualname__
'instrument.<locals>.wrapper'

The changes to these attributes make debugging harder, as it can be hard to tell what functions are referred to. They can also break tools that rely on introspection. “Transparent” decorators preserve these attributes so that the full details of the source function can be used in debugging and introspection.

Copying Attributes for Transparency

We can make our decorator a little bit more transparent by making it copy over __qualname__:

import time


def instrument(wrapped):
    def wrapper(*args, **kwargs):
        start = time.time_ns()
        result = wrapped(*args, **kwargs)
        end = time.time_ns()
        print(f"Took {end - start} nanoseconds")
        return result

    wrapper.__qualname__ = wrapped.__qualname__

    return wrapper

Now the wrapped function’s representation includes the correct name:

>>> @instrument
... def render_html(title):
...     return f"<h1>{title}</h1>"
...

>>> render_html.__qualname__
'render_html'

>>> render_html
<function __main__.render_html(*args, **kwargs)>

The argument specification is still wrong though, as well as several other attributes we haven’t looked at. We can copy all of them by using the standard library function functoools.update_wrapper(). We would use it in our decorator like so:

import functools
import time


def instrument(wrapped):
    def wrapper(*args, **kwargs):
        start = time.time_ns()
        result = wrapped(*args, **kwargs)
        end = time.time_ns()
        print(f"Took {end - start} nanoseconds")
        return result

    functools.update_wrapper(wrapper, wrapped)

    return wrapper

Now when we check the function representation, it’s the same as without the decorator:

>>> @instrument
... def render_html(title):
...     return f"<h1>{title}</h1>"

>>> render_html
<function __main__.render_html(title)>

This is great. If the function representation was all we cared about for transparency, we could stop here. Unfortunately, there are some other differences that might affect certain applications.

Prefect Transparency with Wrapt

The author of the wrapt library that we use, Graham Dumpleton, wrote a series of blog posts on issues with decorators. It starts with the post How you implemented your Python decorator is wrong, which identifies these issues:

We solved the first two by using functools.update_wrapper(), as it copies over the relevant attributes. However, the second two are much more difficult to solve, and we’re glad wrapt does them for us. They’re a bit too much to describe here, but if you’re interested, you should check out the blog post series.

We can reimplement our above decorator using wrapt like so:

import time
import wrapt


@wrapt.decorator
def instrument(wrapped, instance, args, kwargs):
    start = time.time_ns()
    result = wrapped(*args, **kwargs)
    end = time.time_ns()
    print(f"Took {end - start} nanoseconds")
    return result

wrapt.decorator is a decorator for creating decorators. It takes our instrument() function, and turns it into a fully-featured transparent decorator. It saves us writing two levels of function, allowing us to focus solely on our wrapper.

Our instrument function takes four arguments that wrapt passes us on each call:

The body of instrument is similar to the previous implementation, with wrapped(*args, **kwargs) calling the wrapped function.

The wrapped function works exactly the same as our previous implementation, and has the same, transparent representation:

>>> render_html("How to Fetch")
Took 2000 nanoseconds
'<h1>How to Fetch</h1>'

>>> render_html
<function __main__.render_html(title)>

And That’s a Wrapt

wrapt also comes with some other neat features:

To learn more, check out its documentation, and the series of blog posts.

You may also be interested to see our usage in the Scout Python integration. A good starting point is the instrumentation of Jinja2’s Template.render(). There’s also the pull request where we moved to wrapt.

Fin

May your code be well decorated,

—Adam


Working on a Django project? Check out my book Speed Up Your Django Tests which covers loads of best practices so you can write faster, more accurate tests.


Subscribe via RSS, Twitter, or email:

One summary email a week, no spam, I pinky promise.

Related posts:

Tags: python