Python Tutorial: Python Decorators

 ·  · 

Explanation of how Python decorators work.

 

1 Prerequisite

I assume that you feel comfortable with following features of Python functions before moving on. Otherwise, please read my tutorial about Python functions and closures.

  1. Python functions can be passed as arguments to a function
  2. Python functions can be returned as the result of a function

In short, Python functions are higher-order functions.

If you are familiar with closure, so much the better. Otherwise, doesn't matter though.

All examples in this article use Python 3.3+.

 

2 Why decorator?

Suppose we have following transaction functions in our project.

def deposit(money):
    invoice = "Invoice: deposit ${}".format(money)
    return invoice

def transfer(money, dest):
    invoice = "Invoice: transfer ${} to account {}".format(money, dest)
    return invoice

def withdraw(money):
    invoice = "Invoice: withdraw ${}".format(money)
    return invoice

# ......

There is a new requirement now: log the transaction calls but do not change the source codes of original transaction functions.

Firstly, a straightforward solution is to create a wrapper function and print log in the wrapper. Take the deposit function as an example.

def log_transaction(a_func):
    def wrapper(money):
        print('log => calling {}...'.format(a_func.__name__))
        return a_func(money)
    return wrapper

def deposit(money):
    invoice = "Invoice: deposit ${}".format(money)
    return invoice

# ......

Note that log_transaction will return a wrapper function of deposit. This wrapper function is responsible for printing log, then just call the original deposit function and return the result.

Let's try the solution according to steps we've discussed.

  1. obtain the wrapper function
  2. assign the returned wrapper function to the original function variable
  3. call the wrapped function as usual
deposit = log_transaction(deposit)

# Now, deposit is wrapped, but call it as before
invoice = deposit(1000)
print(invoice)

# output:
# log => calling deposit...
# Invoice: deposit $1000

Fortunately Python provide decorator grammar @ (@decorator) to simplify tasks like above.

Now refactor above example using Python decorator.

def log_transaction(a_func):
    def wrapper(money):
        print('log => calling {}...'.format(a_func.__name__))
        return a_func(money)
    return wrapper

# Use "log_transaction" decorator to decorate "deposit"
@log_transaction
def deposit(money):
    invoice = "Invoice: deposit ${}".format(money)
    return invoice

# ......

# No need to reassign
invoice = deposit(999)
print(invoice)

# output:
# log => calling deposit...
# Invoice: deposit $999

 

3 Improve the decorator

3.1 Support flexible number of arguments

Now the log_transaction decorator can work fine with deposit or withdraw which take only one argument in function signature.

Next try to use it to decorate transfer function.

def log_transaction(a_func):
    def wrapper(money):
        print('log => calling {}...'.format(a_func.__name__))
        return a_func(money)
    return wrapper

@log_transaction
def transfer(money, dest):
    invoice = "Invoice: transfer ${} to account {}".format(money, dest)
    return invoice

# ......

invoice = transfer(600, "Jack")
print(invoice)

However, a TypeError raised.

Traceback (most recent call last):
    ......
    transfer(600, "Jack")
TypeError: wrapper() takes 1 positional argument but 2 were given

Because the wrapper function in above log_transaction decorator only accepts one argument, which is the same as deposit or withdraw.

We can make use of Python's arguments unpacking to improve the implementation of wrapper.

def log_transaction(a_func):
    def wrapper(*args, **kwargs):
        print('log => calling {}...'.format(a_func.__name__))
        return a_func(*args, **kwargs)
    return wrapper

So log_transaction decorator can decorate any functions now.

3.2 Attributes of decorated function

First of all, try to inspect __name__ and __doc__ attributes of deposit function after using log_transaction decorator.

def log_transaction(a_func):
    def wrapper(*args, **kwargs):
        """Docstring of wrapper"""
        print('log => calling {}...'.format(a_func.__name__))
        return a_func(*args, **kwargs)
    return wrapper

@log_transaction    
def deposit(money):
    """Docstring of deposit"""
    invoice = "Invoice: deposit ${}".format(money)
    return invoice


print(deposit.__name__)
print(deposit.__doc__)

The printing results will like as follows.

wrapper
Docstring of wrapper

So some attributes of a function object will be modified when the function is wrapped by decorators. Because the decorated deposit is indeed the wrapper.

Intuitively, you may think that we can change values of attributes back in the wrapper function.

def log_transaction(a_func):
    def wrapper(*args, **kwargs):
        """Docstring of wrapper"""
        print('log => calling {}...'.format(a_func.__name__))
        return a_func(*args, **kwargs)

    wrapper.__name__ = a_func.__name__
    wrapper.__doc__ = a_func.__doc__
    return wrapper

Yes, this will work. Obviously, it will be very inconvenient when we write a new decorator in this way. And it is so boring to reassign each attributes of original function to the new wrapper.

But, Python will save your life! Python provided a module named functools which will complete reassignment of attributes automatically.

import functools

def log_transaction(a_func):
    # Magic is here
    @functools.wraps(a_func)
    def wrapper(*args, **kwargs):
        """Docstring of wrapper"""
        print('log => calling {}...'.format(a_func.__name__))
        return a_func(*args, **kwargs)
    return wrapper

@log_transaction    
def deposit(money):
    """Docstring of deposit"""
    invoice = "Invoice: deposit ${}".format(money)
    return invoice


print(deposit.__name__)
print(deposit.__doc__)

# output:
# deposit
# Docstring of deposit

Attributes came back!

 

4 Decorator with arguments

More practically, you may be asked to implement an option of printing custom log level, such as debug, info, warn, error.

What we want is something like following codes.

@log_transaction('DEBUG')
def deposit(money):
    pass

@log_transaction('WARN')
def transfer(money, dest):
    pass

It is only natural that Python decorator can take arguments. Because Python decorator essentially is a Python function. Let's analyze how to implement decorators which have arguments.

@log_transaction('DEBUG')
def deposit(money):
    pass

We know that above decorator grammar is equivalent to following function calls.

deposit = log_transaction('DEBUG')(deposit)

Specifically, split it to steps below.

actual_log_decorator = log_transaction('DEBUG')
deposit = actual_log_decorator(deposit)

or

actual_log_decorator = log_transaction('DEBUG')

@actual_log_decorator
def deposit(money):
    pass

OK, I think it's clear enough to refactor the log_transaction now.

import functools

def log_transaction(level):
    def actual_log_decorator(a_func):
        @functools.wraps(a_func)
        def wrapper(*args, **kwargs):
            print('[{}] => calling {}...'.format(level, a_func.__name__))
            return a_func(*args, **kwargs)
        return wrapper
    return actual_log_decorator

@log_transaction('WARN')
def deposit(money):
    result = "Invoice: deposit ${}".format(money)
    return result

@log_transaction('DEBUG')
def transfer(money, target):
    result = "Invoice: transfer ${} to {}".format(money, target)
    return result

print(deposit(1000))
print(transfer(600, "Jack"))

Run above snippets and the printing results should be as follows.

[WARN] => calling deposit...
Invoice: deposit $1000
[DEBUG] => calling transfer...
Invoice: transfer $600 to Jack

Voila! Our decorator with argument works fine.