Python Tutorial: Python Decorators That Takes Optional Arguments

 ·  · 

Introduce Python partial functions and demonstrate how to use partial functions to define a decorator that takes optional arguments.

 

1 Prerequisite

Basic knowledge of Python functions and Python decorators if you need.

  1. Python Function And Python Closure
  2. Python Decorators

All examples in this article use Python 3.3+.

 

2 Python partial functions

2.1 What is a partial function?

Firstly given a function takes multiple arguments. Following is a simple example function to sell air tickets. It takes three arguments: how many tickets people want to buy, departure and destination.

def tickets_seller(pieces, departure, destination):
    print("Thanks, you bought {} tickets from {} to {}".format(pieces, departure, destination))

Now consider this case: there is a small ticket agency which only and always sells the tickets from New York to Boston.

This means that the arguments departure and destination will be filled respectively with New York and Boston every time the tickets_seller is called. In such a case, it's better to create a convenient function with latter two arguments filled in.

def tickets_seller(pieces, departure, destination):
    print("Thanks, you bought {} tickets from {} to {}".format(pieces, departure, destination))

def shortcut_tickets_seller(pieces):
    return tickets_seller(pieces, 'New York', 'Boston')

shortcut_tickets_seller(2)

# output:
# Thanks, you bought 2 tickets from New York to Boston

The convenient function shortcut_tickets_seller defined here can be called a partial function of tickets_seller.

2.2 Python partial function toolkit

In fact you can create a partial function handily using Python's functools.partial function (in functools module) instead of declaring a new function by def.

import functools

def tickets_seller(pieces, departure, destination):
    print("Thanks, you bought {} tickets from {} to {}".format(pieces, departure, destination))

shortcut_tickets_seller = functools.partial(tickets_seller, departure='New York', destination='Boston')

shortcut_tickets_seller(2)
shortcut_tickets_seller(5)

To sum up, partial function actually is new version of a function with one or more arguments already filled in.

 

3 Define a Python decorator that takes optional arguments

3.1 Review

In the first place, let's review a decorator that takes required arguments.

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('DEBUG')
def deposit(money):
    return "Invoice: deposit $()".format(money)

@log_transaction('WARN')
def withdraw(money):
    return "Invoice: withdraw $()".format(money)

if __name__ == "__main__":
    print(deposit(100))
    print(withdraw(99))

3.2 Goal

Our goal is to refactor above codes to make the level argument optional. In other words, both following ways will work.

# No argument
@log_transaction
def deposit(money):
    return "Invoice: deposit $()".format(money)

# Pass argument
@log_transaction('WARN')
def withdraw(money):
    return "Invoice: withdraw $()".format(money)

3.3 Analysis

So what's the difference in essence between the two ways? If you really understand how a Python decorator works, the answer is easy in fact. Otherwise, it doesn't matter, I will show you the difference at once.

import functools

def log_transaction(level):
    # print values of arguments passed in
    print(locals())

    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
def deposit(money):
    return "Invoice: deposit $()".format(money)

@log_transaction('WARN')
def withdraw(money):
    return "Invoice: withdraw $()".format(money)

if __name__ == "__main__":
    print(deposit(100))
    print(withdraw(99))

Note that I added a line of print(locals()) statement to inspect the arguments passed in.

Run above code snippets and there will be two extra output lines printed by print(locals()).

{'level': <function deposit at 0x10cbc7ea0>}
{'level': 'WARN'}

The first line is generated by following codes block.

@log_transaction
def deposit(money):
    return "Invoice: deposit $()".format(money)

So the @log_transaction, which is without explicit argument, passed the target deposit function to log_transaction function — of course, this is how the Python decorators work: take the target function as argument, wrap it and return the wrapped function — Nothing is special here.

However, according to the second output line, passed argument is only level with its value WARN in the second @log_transaction('WARN') case.

The tricky part is exactly here. As for @log_transaction('WARN'), the actual decorator function is the function returned by log_transaction('WARN') rather than log_transaction itself.

Obviously, the log_transaction function should handle two circumstances so as to achieve our goal.

  1. If the target function is passed in, just work as usual, return the wrapped result function;
  2. If the target function is not passed in, return log_transaction('WARN') first.

We can implement our goal now base on the analysis above.

3.4 Solution

It's easy to give a straightforward solution.

import functools

def log_transaction(a_func=None, level='DEBUG'):
    # if a_func is callable, return wrapped function
    if callable(a_func):
        @functools.wraps(a_func)
        def wrapper(*args, **kwargs):
            print('[{}] => calling {}...'.format(level, a_func.__name__))
            return a_func(*args, **kwargs)

        return wrapper

    # if a_func is not callable, return log_transaction with the second parameter filled in 
    else:
        def partial_log_transaction(a_func):
            return log_transaction(a_func, level=level)

        return partial_log_transaction

# Argument level is optional
@log_transaction
def deposit(money):
    return "Invoice: deposit ${}".format(money)

# Pass keyword argument: level
@log_transaction(level='WARN')
def withdraw(money):
    return "Invoice: withdraw ${}".format(money)

if __name__ == "__main__":
    print(deposit(100))
    print(withdraw(99))

Run it to see output.

[DEBUG] => calling deposit...
Invoice: deposit $100
[WARN] => calling withdraw...
Invoice: withdraw $99

Great, it works well! Remember we have learned how to use Python partial functions. So we can replace definition of partial_log_transaction with a functools.partial call.

import functools

def log_transaction(a_func=None, level='DEBUG'):
    # if a_func is callable, return wrapped function
    if callable(a_func):
        @functools.wraps(a_func)
        def wrapper(*args, **kwargs):
            print('[{}] => calling {}...'.format(level, a_func.__name__))
            return a_func(*args, **kwargs)

        return wrapper

    # if a_func is not callable, return log_transaction with the second parameter filled in 
    else:            
        return functools.partial(log_transaction, level=level)

That's it.