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?

Disordered

@KitHo -- it's a boolean flag, so using a boolean value is the right approach. – Daren

@KitHo -- what is "gd"? Is it "good"? – Chimaera

Does this answer your question? Decorators with parameters? – Saleswoman

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

Humidity

So basically decorator always takes only one argument which is the function. But the decorator can be a return value of a function which might take arguments. Is this correct? – Disordered

@balki: Yes, that's correct. What confuses things is that many people will also call the outer function (myDecorator here) a decorator. This is convenient for a user of the decorator, but can be confusing when you're trying to write one. – Humidity

Small detail that confused me: if your log_decorator takes a default argument, you cannot use @log_decorator, it must be @log_decorator() – Nygaard

What if I don't want to pass True to @log_decorator? I want to pass the parameter that the function 'f' gets. – Martainn

This solution is very elegant. The decorator returns a decorator which returns a wrapper which returns a function – Pascal

@RonyTesler The parameters of the function 'f' are passed to the wrapper within your actual decorator (`def wrapper(*args, **kwargs):' above, but you can also name specific arguments that f will receive). If there are no parameters you want to pass to @log_decorator at definition time, you don't even need it; just define 'actual_decorator' and put '@actual_decorator' (no parentheses!) before your function definition. – Huntley

Answer 2

Just to provide a different viewpoint: the syntax

@expr
def func(...): #stuff

is equivalent to

def func(...): #stuff
func = expr(func)

In particular, expr can be anything you like, as long as it evaluates to a callable. In particular particular, expr can be a decorator factory: you give it some parameters and it gives you a decorator. So maybe a better way to understand your situation is as

dec = decorator_factory(*args)
@dec
def func(...):

which can then be shortened to

@decorator_factory(*args)
def func(...):

Of course, since it looks like decorator_factory is a decorator, people tend to name it to reflect that. Which can be confusing when you try to follow the levels of indirection.

Pollack

Thanks, this really helped me understand the rationale behind what's going on. – Mostly

Answer 3

Just another way of doing decorators. I find this way the easiest to wrap my head around.

class NiceDecorator:
    def __init__(self, param_foo='a', param_bar='b'):
        self.param_foo = param_foo
        self.param_bar = param_bar

    def __call__(self, func):
        def my_logic(*args, **kwargs):
            # whatever logic your decorator is supposed to implement goes in here
            print('pre action baz')
            print(self.param_bar)
            # including the call to the decorated function (if you want to do that)
            result = func(*args, **kwargs)
            print('post action beep')
            return result

        return my_logic

# usage example from here on
@NiceDecorator(param_bar='baaar')
def example():
    print('example yay')


example()

Prostitute

thank you! been looking at some mind-bending "solutions" for about 30 minutes and this is the first one that actually makes sense. – Piggott

Pretty clever way to achieve decorator with optional params, without the need to create nesting complexity. Refactoring all my custom decorators to this, thanks! – Bedspring

The only down side is that it cannot be used without explicitly calling the decorator when you just want to use the default parameters. You have to decorate like @NiceDecorator() even when you are not passing any params. I wonder if it can be modified to enable writing @NiceDecorator when passing no params? – Mercurialize

Answer 4

Just want to add some usefull trick that will allow to make decorator arguments optional. It will also alows to reuse decorator and decrease nesting

import functools

def myDecorator(test_func=None,logIt=None):
    if test_func is None:
        return functools.partial(myDecorator, logIt=logIt)
    @functools.wraps(test_func)
    def f(*args, **kwargs):
        if logIt==1:
            print 'Logging level 1 for {}'.format(test_func.__name__)
        if logIt==2:
            print 'Logging level 2 for {}'.format(test_func.__name__)
        return test_func(*args, **kwargs)
    return f

#new decorator 
myDecorator_2 = myDecorator(logIt=2)

@myDecorator(logIt=2)
def pow2(i):
    return i**2

@myDecorator
def pow3(i):
    return i**3

@myDecorator_2
def pow4(i):
    return i**4

print pow2(2)
print pow3(2)
print pow4(2)

Phidippides

Sorry, for the necro bump. I'm interested if this pattern is considered a good practice? – Gerrald

Answer 5

Now if you want to call a function function1 with a decorator decorator_with_arg and in this case both the function and the decorator take arguments,

def function1(a, b):
    print (a, b)

decorator_with_arg(10)(function1)(1, 2)

Litchfield

Answer 6

Decorator that accept multiple input is like dataclasses decorator

In that example, dataclass accept three syntaxes:

@dataclass
class C:
    ...

@dataclass()
class C:
    ...

@dataclass(init=True, repr=True, eq=True, order=False, unsafe_hash=False, frozen=False,
           match_args=True, kw_only=False, slots=False, weakref_slot=False)
class C:
    ...

In creating decorator with same behavior, you can use this,

import inspect

def customdecorator(*args, **kwargs):
    def decorator(func):
        print('input decorator:', args, kwargs)
        def __wrapper(*func_args, **func_kwargs):
            print('input decorator inside function:', args, kwargs)
            print('input function:', func_args, func_kwargs)
            # Do something before calling the function
            result = func(*func_args, **func_kwargs)
            # Do something after calling the function
            return result
        return __wrapper

    print('input root:', args, kwargs)
    if len(kwargs) > 0:
        # Decorator is used with arguments, e.g., @functionmethod(arg1=val1, arg2=val2)
        return decorator
    if len(args) == 0:
        return decorator
    if len(args) == 1:
        return decorator(args[0])

# Example usages
@customdecorator
def example1():
    print("Function without call")

@customdecorator()
def example2():
    print("Function without arguments")

@customdecorator(arg1="value1", arg2="value2")
def example3(arg2):
    print(f"Function with arguments: {arg2}")

example1()
example2()
example3(arg2="ex2")

In your case, it will be

def myDecorator(*args, **kwargs):
    def decorator(func):
        def __wrapper(*func_args, **func_kwargs):
            # Do something before calling the function
            test_func = kwargs.get('test_func', None)
            logIt = kwargs.get('logIt', None)
            if logIt:
                print("Calling Function: " + test_func.__name__)

            result = func(*func_args, **func_kwargs)

            # Do something after calling the function
            return result
        return __wrapper

    if len(kwargs) > 0:
        # Decorator is used with arguments, e.g., @functionmethod(arg1=val1, arg2=val2)
        return decorator
    if len(args) == 0:
        return decorator
    if len(args) == 1:
        return decorator(args[0])


@myDecorator(logIt=False)
def someFunc():
    print('Hello')

someFunc()

The caveat is:

1. Optional parameter must use key argument.
2. Default value is entered via kwargs.get(..., ...).

Rarotonga