Python: Decorators with Arguments

A decorator in Python in any callable object that modifies another existing function or class.

Because a decorator is callable, it can also accept arguments.

@decorator(arg1, arg2)
def function():
  ...

In this tutorial, we will learn to write custom decorators that can also accept additional arguments.

Function Decorator with Arguments

The following is an example of a custom decorator that doesn’t accept any arguments.

#custom decorator that does not accept arguments
def mydecorator(func):
    def newfunc(num):
        print("calling func from decorator")
        func(num)
    return newfunc
    
@mydecorator
def myfunc(num):
    print(num*num)
    
myfunc(7)

Output:

calling func from decorator
49

To make this decorator accept arguments, we have to wrap it with another function (that accepts arguments to decorator) as follows:

def argdecorator(arg1, arg2):
    print('Arguments: {}, {}'.format(arg1, arg2))
    def mydecorator(func):
        def newfunc(num):
            print("calling func from decorator")
            func(num)
        return newfunc
    return mydecorator

Now if we apply this decorator to a function, we would be able to pass arguments.

@argdecorator('hello', 'world')
def myfunc(num):
    print(num*num)
    
myfunc(7)

Output:

Arguments: hello, world
calling func from decorator
49

Class Decorator with Arguments

In custom class decorator, when there is no arguments to the decorator, the reference of the implicit function is passed to the __init__ method.

class Decorator:
    def __init__(self, func):
        self.func = func
    
    def __call__(self, num):
        print('calling func from class decorator')
        self.func(num)

@Decorator
def myfunc(num):
    print(num*num)
    
myfunc(5)

In this case, the __call__ method works as the decorator method that modifies the existing/passed function.

If we pass arguments to the class decorator, they (in place of the implicit function) are passed to the __init__ method.

In such case, we accept the reference of the implicit function in the declaration of the __call__ method, and write it as the function decorator.

class Decorator:
    def __init__(self, arg1, arg2):
        self.arg1 = arg1
        self.arg2 = arg2
    
    def __call__(self, func):
        print('Arguments: {}, {}'.format(self.arg1, self.arg2))
        def mydecorator(num):
            print('calling func from class decorator')
            func(num)
        return mydecorator

@Decorator('hello', 'world')
def myfunc(num):
    print(num*num)
    
myfunc(5)

Output:

Arguments: hello, world
calling func from class decorator
25