Writing Python Decorators

Writing Python Decorators

You might think it’s tricky to write decorators in Python. Well, this article is written to demistify that process.

Decorators are simply functions that allow replacing and/or wrapping functions with other code, in a syntax friendly format when using the decorator. You can use them to ensure a piece of code executes around other function calls, or any other cross cutting concern you might have (transactions, logging, and security come immediately to mind).

For example in Germanium we use the @iframe decorator to ensure that we are in a certain IFrame when executing commands.

Flask (another popular REST framework) uses this for configuring its endpoints.

Writing Simple Decorators

In order to make something somewhat "usable" we will write a decorator that logs function executions.

def trace(f):  # (1)
    def wrapper(*args, **kw):  # (2)
        try:
            print("=> entering")
            return f(*args, **kw)
        finally:
            print("<= leaving")

    return wrapper

1. We define a decorator with the name @trace. Actually any function can be a decorator, and all it needs to do is to return another callable that will replace the decorated callable.

2. The returned function is going to replace the original function.

If we would be to apply it we can just:

@trace
def add(a, b):
    print("computing %d + %d" % (a, b))
    return a + b

This code is actually equal to:

def add(a, b):
    print("computing %d + %d" % (a, b))
    return a + b

add = trace(add)

It’s just prettier.

As you might have guessed, running:

print(add(1,2))

Will output the following on the screen:

=> entering
computing 1 + 2
<= leaving
3

Writing Configurable Decorators

The decorators themselves can also be called when applied:

@trace("addition")
def add(a, b):
    print("computing %d + %d" % (a, b))
    return a + b

but writing the decorator must be done slightly differently:

def trace(name): # (1)
    def decorator(f): # (2)
        def wrapper(*args, **kw):
            try:
                print("=> entering %s" % name) # (3)
                return f(*args, **kw)
            finally:
                print("<= leaving %s" % name)

        return wrapper
    return decorator

This time the trace function instead of returning directly the wrapping function function itself, will return the decorator function instead. This is because the result of evaluating a decorator must be a function reference that will receive the function to wrap, and will return the wrapper to do the actual work.

Again, this would actually be just:

def add(a, b):
    print("computing %d + %d" % (a, b))
    return a + b

add = trace("addition")(add)

Chaining Decorators

One of the coolest things in the decorator land from python is that you can actually chain the decorators. That means a declarations like this are possible:

@trace
@transactional
def run_some_code():
    # ...

The order of the decorators now becomes important, but if you were paying attention, you know that when calling the run_some_code() the actual code executed will probably be:

-> trace
  -> transactional
    -> run_some_code

This is why I said that decorators are great for using them, because they clarify intent. Compare with:

def run_some_code():
    # ...
run_some_code = trace(transactional(run_some_code))

Pretty Decorators

Since every decorator returns a "wrapper" function, that wraps the original one, you might wonder what happens when you want to see what function is actually there?

If in our code sample you would do:

print(run_some_code)

You might be surprised to see on the screen:

<function trace.<locals>.decorator.<locals>.wrapper at 0x7f9bf808eae8>

If you we’re paying attention, you know now why, since this is just a wrapped function. Either way it’s not really helpful to see what function is actually getting called, since most decorators by design are non-functional.

The good news is that python provides in the functools package a function named wraps (itself a decorator) that can augument the function we return with the original function name that we’re wrapping, and also its original documentation.

Having the import on wraps from functools allows us to add the @wraps call:

def trace(name):
    def decorator(f):
        @wraps(f)
        def wrapper(*args, **kw):
            try:
                print("=> entering %s" % name)
                return f(*args, **kw)
            finally:
                print("<= leaving %s" % name)

        return wrapper
    return decorator

This is why when printing it now, you’ll just see:

<function run_some_code at 0x7eff83a7ea60>

Conclusions

1. Decorators are just functions that receive themselves a function, so they can change the behavior, including even returning a completely different function. Most of the time, as seen, they just wrap the original function with some other call.

2. Decorators notation supports also calling, but then you need to create a function that creates the decorator, that will wrap the function.

3. Use functools.wraps to copy the original name into the returned function from the decorator.

4. Here is a full code sample.