Decorator

Decorators in Python are functions that modify other functions. If we need to modify a lot of functions all at once with a minimum of coding, we should use a decorator. A decorator allows us to modify the behavior of a function without changing the codes inside of the function permanently.

A decorator is defined just like a function, but it must take another function as input argument. The decorator at the end must return the modified function. To modify the function, the typical way to this is to define the new function inside the decorator. The new function inside the decorator is sometimes called a “wrapper”.

Suppose we have made a bunch of one-variable functions that simply “return” some output. But somehow later we actually want those functions to “print” the output instead. In this case, we can make just make a decorator that modifies functions to accomplish this for all functions at once instead of changing the codes in every function.

Suppose the functions that we made look like this:

# function that takes one variable and return some output
def my_func(x):
    return (x + 3) ** 2

And a decorator that prints the output instead of returning:

# a decorator definition
def my_dec(fun):
    def wrapper(x):
        print(fun(x))
    return wrapper

Notice that a decorator must take a function as input argument, and define a new function that modifies that original function. In the end, the decorator must return the name of the function (NOT INCLUDE any arguments). What’s more, the decorator must be above the function that is being applied.

To apply a decorator to a function definition, simply put @<name of decorator> right above the “def”. When we call the original function, the behavior has been changed.

@my_dec
def my_func(x):
    return (x + 3) ** 2

my_func(3)

Output:

36

The number of input argument of new function inside a decorator is not necessarily match the original function. Instead, when we call the function that has decorator, the input arguments need to match the new function inside the decorator. For example:

# a wrapper that takes no input argument
def take_input(fun):
    def wrapper():
        x = eval(input("Please enter a number: "))
        print(fun(x))
    return wrapper

# a function take takes one input argument
@take_input
def my_func(x):
    return (x + 3) ** 2


my_func()

Output:

Please enter a number: 5
64

We can apply multiple decorators to a function. But we need to pay attention to the order of decorators (ORDER MATTERS).

# a decorator that prints some logs
def logger(fun):
    def wrapper():
        print("before the function")
        fun()
        print("after the function")
    return wrapper

# a decorator that runs the function twice
def do_twice(fun):
    def wrapper():
        fun()
        fun()
    return wrapper

# the closer decorator will be applied earlier
@logger
@do_twice
def new_func():
    print("hello world")

Output:

before the function
hello world
hello world
after the function

To make our decorator as flexible as possible to handler all inputs, use *args and **kwargs to allow for arbitrarily many arguments and keyword arguments.

def error_handler(fun):
    def wrapper(*args, **kwargs):
        try:
            return fun(*args, **kwargs)
        except Exception as e:
            return f"An error occurs! The error is {e}"
    return wrapper


@error_handler
def division(a, b):
    return a / b

Output:

2.5
An error occurs! The error is division by zero

Decorators can be used for all sort of things, like graphing or debugging or timing your functions or suppressing/handling errors or printing output, etc. Another common used of a decorator is timing our functions. For example:

def time_dec(fun):
    def wrapper(*args, **kwargs):
        start = time.time()
        out = fun(*args, **kwargs)
        end = time.time()
        print(f"It took {end - start} seconds for {fun.__name__} to finish")
        return out
    return wrapper


@time_dec
def slow_fun(n):
    ret = 0
    for i in range(n ** n):
        ret += 1
        ret -= 1
    return ret


slow_fun(9)

Output:

It took 11.597818851470947 seconds for slow_fun to finish