File size: 3,917 Bytes
d8d14f1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
from functools import wraps
from time import time
from typing import Any, Callable

from swarms.utils.loguru_logger import initialize_logger

logger = initialize_logger("try_except_wrapper")


def retry(
    max_retries: int = 3,
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """
    A decorator that retries a function a specified number of times if an exception occurs.

    Args:
        max_retries (int): The maximum number of retries. Default is 3.

    Returns:
        Callable[[Callable[..., Any]], Callable[..., Any]]: The decorator function.
    """

    def decorator_retry(
        func: Callable[..., Any]
    ) -> Callable[..., Any]:
        @wraps(func)
        def wrapper_retry(*args, **kwargs) -> Any:
            """
            The wrapper function that retries the decorated function.

            Args:
                *args: Variable length argument list.
                **kwargs: Arbitrary keyword arguments.

            Returns:
                Any: The result of the decorated function.
            """
            for _ in range(max_retries):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    logger.error(f"Error: {e}, retrying...")
            return func(*args, **kwargs)

        return wrapper_retry

    return decorator_retry


def log_execution_time(
    func: Callable[..., Any]
) -> Callable[..., Any]:
    """
    A decorator that logs the execution time of a function.

    Args:
        func (Callable[..., Any]): The function to be decorated.

    Returns:
        Callable[..., Any]: The decorated function.
    """

    @wraps(func)
    def wrapper(*args, **kwargs) -> Any:
        """
        The wrapper function that logs the execution time and calls the decorated function.

        Args:
            *args: Variable length argument list.
            **kwargs: Arbitrary keyword arguments.

        Returns:
            Any: The result of the decorated function.
        """
        start = time()
        result = func(*args, **kwargs)
        end = time()
        logger.info(
            f"Execution time for {func.__name__}: {end - start} seconds"
        )
        return result

    return wrapper


def try_except_wrapper(verbose: bool = False):
    """
    A decorator that wraps a function with a try-except block.
    It catches any exception that occurs during the execution of the function,
    prints an error message, and returns None.
    It also prints a message indicating the exit of the function.

    Args:
        func (function): The function to be wrapped.

    Returns:
        function: The wrapped function.

    Examples:
    >>> @try_except_wrapper(verbose=True)
    ... def divide(a, b):
    ...     return a / b
    >>> divide(1, 0)
    An error occurred in function divide: division by zero
    Exiting function: divide
    """

    def decorator(func: Callable[..., Any]):
        @wraps(func)
        @retry()
        @log_execution_time
        def wrapper(*args, **kwargs):
            try:
                result = func(*args, **kwargs)
                return result
            except Exception as error:
                if verbose:
                    logger.error(
                        f"An error occurred in function {func.__name__}:"
                        f" {error}"
                    )
                else:
                    logger.error(
                        f"An error occurred in function {func.__name__}:"
                        f" {error}"
                    )
                    return None
            finally:
                print(f"Exiting function: {func.__name__}")

        return wrapper

    return decorator


# @try_except_wrapper(verbose=True)
# def divide(a, b):
#     """Multiply two numbers."""
#     return a / b


# # This will work fine
# result = divide(2, 0)
# print(result)  # Output: 6