Spaces:
Sleeping
Sleeping
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
|