Skip to content

A Custom Decorator For Retry In Python

Posted on:May 17, 2024 at 04:16 AM

tp.web.random_picture

Custom Decorator For A Retry Mechanism For Better Error Handling And Clean Code

Table of contents

Open Table of contents

TL;DR

The retry decorator can be used to retry a function for a specified number of times in case of failure. This is useful when calling an API that might fail due to network issues.

import time

def retry(max_retries=3, delay=1):
    def decorator(original_function):
        def wrapper_func(*args, **kwargs):
 retries = 0
            if retries < max_retries:
                try:
                    return original_function(*args, **kwargs)
                except Exception as e:
                    print(f"Function {original_function.__name__} failed with error {e}. Retrying...")
 retries += 1
 time.sleep(delay)
            else:
                raise Exception(f"Function {original_function.__name__} failed after {max_retries} retries")
        return wrapper_func
    return decorator

@retry(max_retries=3, delay=1)
def my_function():
    print("Hello World")
    raise Exception("Something went wrong")

my_function()

# Function my_function failed with error Something went wrong. Retrying...
# Function my_function failed with error Something went wrong. Retrying...
# Function my_function failed with error Something went wrong. Retrying...
# Traceback (most recent call last):
#   File "<stdin>", line 1, in <module>
#   File "<stdin>", line 13, in wrapper_func
# Exception: Function my_function failed after 3 retries

Introduction

While implementing anything In any programming language it is essential to implement error handling. In Python, we use try...except to handle exceptions. However, there might be a use case where instead of directly handling the exception it’s required to retry the whole execution. Mainly in cases of API failure due to connectivity or rate limit, we can have a retry.

To implement a reusable retry mechanism in Python we can implement a decorator. decorator is a special type of function that can be used to modify the behaviour of another function. This is because Python supports passing a function a parameter in another function

Decorator function basics

A decorator function in Python is a higher-order function that takes a function as an argument and returns another function. The decorator function can modify the behaviour of the function passed to it.

def my_decorator_function(original_function):
    def wrapper_function(*args, **kwargs):
        # Do something before the function is called
        print(f"Function {original_function.__name__} called with args {args} and kwargs {kwargs}")
 result = original_function(*args, **kwargs)
        print(f"Function {original_function.__name__} returned {result}")
        # Do something after the function is called
        return result
    return wrapper_function

In the above example:

To use the decorator function we can use the @ symbol before the function definition

@my_decorator_function
def my_function():
    print("Hello World")

When we can my_function it will be wrapped by the my_decorator_function and the output will be

my_function()
# Function my_function called with args () and kwargs {}
# Hello World
# Function my_function returned None

Real-World Use Cases of Decorator Functions

There are many use cases of decorator functions in Python. Some of the common use cases are

  1. Logging - To log the input and output of a function along with different parameters
  2. Timing - To measure the time a function takes to execute. This is useful for performance optimization
  3. Rate Limiting - To limit the number of times a function can be called in a given time frame
  4. Retry Mechanism - To retry the function in case of failure. This is useful when calling an API that might fail due to network issues.
  5. Authentication - To authenticate the user before calling the function
  6. Caching - To cache the result of a function for future use. This is useful when the function is called multiple times with the same input

Retry Mechanism Using Decorator Function

To implement a retry mechanism using a decorator function we can define a decorator function that takes the number of retries as an argument. The decorator function will call the original function and if it fails it will retry the function for the specified number of times.

import time

def retry(max_retries=3, delay=1):
    def decorator(original_function):
        def wrapper_func(*args, **kwargs):
 retries = 0
            if retries < max_retries:
                try:
                    return original_function(*args, **kwargs)
                except Exception as e:
                    print(f"Function {original_function.__name__} failed with error {e}. Retrying...")
 retries += 1
 time.sleep(delay)
            else:
                raise Exception(f"Function {original_function.__name__} failed after {max_retries} retries")
        return wrapper_func
    return decorator

To use the retry decorator we can use the following syntax

@retry(max_retries=3, delay=1)
def my_function():
    print("Hello World")
    raise Exception("Something went wrong")

my_function()

In the above example, the my_function will be retried 3 times with a delay of 1 second between each retry. If the function fails after 3 retries it will raise an exception. The output of the above code will be

my_function()
# Function my_function failed with error Something went wrong. Retrying...
# Function my_function failed with error Something went wrong. Retrying...
# Function my_function failed with error Something went wrong. Retrying...
# Traceback (most recent call last):
#   File "<stdin>", line 1, in <module>
#   File "<stdin>", line 13, in wrapper_func
# Exception: Function my_function failed after 3 retries

Conclusion

In this article, we learned how to implement a retry mechanism using a decorator function in python. Decorator functions are a powerful feature of python that allows us to modify the behaviour of a function without changing its code. We also learned about some real-world use cases of decorator functions and how they can be used to improve the readability and maintainability of our code.