Beautify your Python code with Decorators

What are Decorators?

Decorators are functions that modify or add additional functionality to the behavior of other functions or classes without changing the source code of the other functions. In Python, decorators can be identified by the “@” symbol followed by the decorator's name. This is placed immediately before the definition of a function or class that is being decorated.

What is the need for Decorators?

Consider the user info endpoint in a website that displays the user’s information. This endpoint must display the information only if the user is logged in. If the user is not logged in, it should redirect to the login page.

One can modify the source code of this endpoint to check if the user is logged in and then decide whether to display the info or redirect. Although it seems like an easy problem to fix, consider that there are many endpoints that must work only if the user is logged in. Changing the source code for each endpoint to check and redirect leads to code duplication.

Decorators provide the best solution to this problem and allow not only code re-usability and maintenance but also various other functionalities. Let’s look at how decorators can be used to solve the above problem and a few other examples where decorators can be used.

What is the syntax for Decorators?

import functools

def decorator_function(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        # do something before the function call
        value = func(*args, **kwargs)
        # do something after the function call
        return value
    
    return wrapper

@decorator_function
def new_func():
    # do something

This is the general syntax for a decorator function. The decorator_function is defined just like any other function except that it takes a function reference that needs to be wrapped as the argument.

It is a good practice to use the functools.wraps decorator to preserve the information about the original function or class, otherwise, information such as the original function’s name, docstring, arguments, etc., will be overwritten with the decorator's information.

A wrapper function that takes any number of positional and keyword arguments is defined that calls the original function with the same arguments. It does some functionality before and/or after calling the original function.

The “@” symbol along with the decorator function name can then be used to wrap any function and provide it with additional functionalities.

What are some examples where decorators can be used?

1. Logging

A simple logging decorator can be created that logs all the arguments and the return value when a function is called.

import functools

def log_function(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        args_repr = [repr(a) for a in args]
        kwargs_repr = [f"{k}={v!r}" for k, v in kwargs.items()]
        signature = ", ".join(args_repr + kwargs_repr)
        
        print(f"Calling function {func.__name__}({signature})")
        
        value = func(*args, **kwargs)
        print(f"function {func.__name__!r} returned {value!r}")
        
        return value
    
    return wrapper

@log_function
def is_prime(n):
    if n < 2: 
         return False
    if n % 2 == 0:
        if n==2:
            return True
        return False
    k = 3
    while k*k <= n:
         if n % k == 0:
             return False
         k += 2
    return True

@log_function
def hello():
    print("Hello World")
    
hello()
print(is_prime(25))

#########################
# Output: 
# 
# Calling function hello()
# Hello World
# function 'hello' returned None
# Calling function is_prime(25)
# function 'is_prime' returned False
# False
#########################

2. Timing

A timing decorator can be created that outputs the time is taken to complete the function call

import functools
import time

def timer(func):
    @functools.wraps(func)
    def wrapper(*args, **kargs):
        start_time = time.perf_counter()
        value = func(*args, **kargs)
        end_time = time.perf_counter()
        run_time = end_time - start_time
        print(f"Finished {func.__name__!r} in {run_time:.4f} secs")
        return value
    
    return wrapper

@timer
def sum_upto(n):
    s = 0
    for i in range(1,n+1):
        s+=i
    print(s)

sum_upto(100000)

#########################
# Output: 
# 
# 5000050000
# Finished 'sum_upto' in 0.0029 secs
#########################

3. Login verification

import functools
from flask import *

app = Flask(__name__)
app.secret_key = "dsajdnsakj"

def require_login(func):
    @functools.wraps(func)
    def wrapper(*args, **kargs):
        if "logged_in" in session:
            return func(*args, **kargs)
        return abort(401)
    
    return wrapper

@app.route('/user_info')
@require_login
def user_info():
    return "User Info"

@app.route('/')
def index():
    return "Index Page"

@app.route('/login')
def login():
    session["logged_in"] = True
    return "Logged in"


@app.route('/logout')
@require_login
def logout():
    session.pop("logged_in")
    return "Logged out"

if __name__== "__main__":
    app.run(host="0.0.0.0", debug = True, port = 8000)