How do I pass extra arguments to a Python decorator?

Question

I have a decorator like below.

def myDecorator(test_func):
    return callSomeWrapper(test_func)
def callSomeWrapper(test_func):
    return test_func
@myDecorator
def someFunc():
    print 'hello'

I want to enhance this decorator to accept another argument like below

def myDecorator(test_func,logIt):
    if logIt:
        print "Calling Function: " + test_func.__name__
    return callSomeWrapper(test_func)
@myDecorator(False)
def someFunc():
    print 'Hello'

But this code gives the error,

TypeError: myDecorator() takes exactly 2 arguments (1 given)

Why is the function not automatically passed? How do I explicitly pass the function to the decorator function?

Answer 1

Since you are calling the decorator like a function, it needs to return another function which is the actual decorator:

def my_decorator(param):
    def actual_decorator(func):
        print("Decorating function {}, with parameter {}".format(func.__name__, param))
        return function_wrapper(func)  # assume we defined a wrapper somewhere
    return actual_decorator

The outer function will be given any arguments you pass explicitly, and should return the inner function. The inner function will be passed the function to decorate, and return the modified function.

Usually you want the decorator to change the function behavior by wrapping it in a wrapper function. Here's an example that optionally adds logging when the function is called:

def log_decorator(log_enabled):
    def actual_decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            if log_enabled:
                print("Calling Function: " + func.__name__)
            return func(*args, **kwargs)
        return wrapper
    return actual_decorator

The functools.wraps call copies things like the name and docstring to the wrapper function, to make it more similar to the original function.

Example usage:

>>> @log_decorator(True)
... def f(x):
...     return x+1
...
>>> f(4)
Calling Function: f
5