End Point: Python decorator basics, part II

This is a continuation of my previous post: Python decorator basics. Here I'll talk about a decorator with optional arguments. Let's say we want to pass an optional argument to the same debug decorator:

def debug(msg=None):
    def actual_decorator(f):
        def wrapper(*args):
            if msg:
                print msg
            return f(*args)
        return wrapper
    return actual_decorator

@debug("Let's multiply!")
def mul(x, y):
    return x*y

Calling mul:

mul(5, 2)
Let's multiply!
10

Excellent. Now let's decorate without a msg and call mul:

@debug
def mul(x, y):
    return x*y

mul(5, 2)
Traceback (most recent call last):
  File "", line 1, in 
TypeError: actual_decorator() takes exactly 1 argument (2 given)

Oh oh. Let's see what happens at time of decoration:

mul = debug(mul)

Hmmm, mul gets passed to debug as it's argument and then the arguments (5, 2) are passed to actual_decorator, since debug returns actual_decorator. To resolve this we need to always call the decorator as a function:

@debug()
def mul(x, y):
    return x*y

mul(5, 2)
10

Assuming that we always expect the msg parameter to be a non-callable, another option would be to check the type of argument passed to the debug decorator:

def debug(msg=None):
    def actual_decorator(f):
        def wrapper(*args):
            if msg:
                print msg
            return f(*args)
        return wrapper
    if callable(msg):
        # debug decorator called without an argument so
        # msg is the function being decorated
        return debug()(msg)
    return actual_decorator

@debug
def mul(x, y):
    return x*y

mul(5, 2)
10

@debug("Let's multiply!")
def mul(x, y):
    return x*y

mul(5, 2)
Let's multiply!
10