Skip to content

Instantly share code, notes, and snippets.

@a-luna
Last active August 17, 2024 22:38
Show Gist options
  • Save a-luna/0a20912a82a8b867e19ab3734fce482a to your computer and use it in GitHub Desktop.
Save a-luna/0a20912a82a8b867e19ab3734fce482a to your computer and use it in GitHub Desktop.
Here are three Python decorators that I use quite often: add timeout functionality to potentially long-running functions, add retry-logic to unreliable functions, and measure the execution time of a function. Also, pytest functions are included as usage examples of each decorator.
"""decorators.log_call"""
import inspect
import logging
from datetime import datetime
from functools import wraps
DT_NAIVE = "%Y-%m-%d %I:%M:%S %p"
class LogCall:
"""Log call signature and execution time of decorated function."""
def __init__(self, logger=None):
self.logger = logger
def __call__(self, func):
if not self.logger:
logging.basicConfig()
self.logger = logging.getLogger(func.__module__)
self.logger.setLevel(logging.INFO)
@wraps(func)
def wrapper(*args, **kwargs):
func_call_args = get_function_call_args(func, *args, **kwargs)
exec_start = datetime.now()
result = func(*args, **kwargs)
exec_finish = datetime.now()
exec_time = format_timedelta_str(exec_finish - exec_start)
exec_start_str = exec_start.strftime(DT_NAIVE)
self.logger.info(f"{exec_start_str} | {func_call_args} | {exec_time}")
return result
def get_function_call_args(func, *args, **kwargs):
"""Return a string containing function name and list of all argument names/values."""
func_args = inspect.signature(func).bind(*args, **kwargs)
func_args.apply_defaults()
func_args_str = ", ".join(f"{arg}={val}" for arg, val in func_args.arguments.items())
return f"{func.__name__}({func_args_str})"
def format_timedelta_str(td):
"""Convert timedelta to an easy-to-read string value."""
(milliseconds, microseconds) = divmod(td.microseconds, 1000)
(minutes, seconds) = divmod(td.seconds, 60)
(hours, minutes) = divmod(minutes, 60)
if td.days > 0:
return f"{td.days}d {hours:.0f}h {minutes:.0f}m {seconds}s"
if hours > 0:
return f"{hours:.0f}h {minutes:.0f}m {seconds}s"
if minutes > 0:
return f"{minutes:.0f}m {seconds}s"
if td.seconds > 0:
return f"{td.seconds}s {milliseconds:.0f}ms"
if milliseconds > 0:
return f"{milliseconds}ms"
return f"{td.microseconds}us"
return wrapper
"""decorators.retry"""
from functools import wraps
from time import sleep
class RetryLimitExceededError(Exception):
"""Custom error raised by retry decorator when max_attempts have failed."""
def __init__(self, func, max_attempts):
message = (
f"Retry limit exceeded! (function: {func.__name__}, max attempts: {max_attempts})"
)
super().__init__(message)
def handle_failed_attempt(func, remaining, ex, delay):
"""Example function that could be supplied to on_failure attribute of retry decorator."""
message = (
f"Function name: {func.__name__}\n"
f"Error: {repr(ex)}\n"
f"{remaining} attempts remaining, retrying in {delay} seconds..."
)
print(message)
def retry(*, max_attempts=2, delay=1, exceptions=(Exception,), on_failure=None):
"""Retry the wrapped function when an exception is raised until max_attempts have failed."""
def decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
for remaining in reversed(range(max_attempts)):
try:
return func(*args, **kwargs)
except exceptions as ex:
if remaining > 0:
if on_failure:
on_failure(func, remaining, ex, delay)
sleep(delay)
else:
raise RetryLimitExceededError(func, max_attempts) from ex
else:
break
return wrapper
return decorator
"""tests.test_log_call"""
import time
import logging
from random import randint
from decorators.log_call import LogCall
logging.basicConfig()
log = logging.getLogger("custom_log")
log.setLevel(logging.INFO)
log.info("logging started")
@LogCall()
def save_values(a, b, c):
pass
@LogCall(log)
def rand_time(min=1, max=3, add_random=False):
time.sleep(randint(min, max))
if add_random:
time.sleep(randint(100, 500) / 1000.0)
def test_default_logger(caplog):
save_values("Aaron", "Charlie", "Ollie")
logger, level, message = caplog.record_tuples[-1]
assert logger == "tests.test_log_call"
assert level == logging.INFO
assert "save_values(a=Aaron, b=Charlie, c=Ollie)" in message
def test_custom_logger(caplog):
rand_time(max=4, add_random=True)
logger, level, message = caplog.record_tuples[-1]
assert logger == "custom_log"
assert level == logging.INFO
assert "rand_time(min=1, max=4, add_random=True)" in message
"""tests.test_retry"""
import time
import pytest
from decorators.retry import retry, RetryLimitExceededError
from decorators.timeout import timeout
@retry(max_attempts=2, delay=1, exceptions=(TimeoutError,))
@timeout(seconds=1)
def retry_with_timeout():
time.sleep(2)
def test_retry_with_timeout():
with pytest.raises(RetryLimitExceededError):
retry_with_timeout()
"""tests.test_timeout"""
import time
import pytest
from decorators.timeout import timeout
@timeout(seconds=1)
def sleep():
time.sleep(2)
def test_timeout():
with pytest.raises(TimeoutError):
sleep()
"""decorators.timeout"""
from functools import wraps
from signal import signal, alarm, SIGALRM
def timeout(*, seconds=3, error_message="Call to function timed out!"):
"""Abort the wrapped function call after the specified number of seconds have elapsed."""
def _handle_timeout(signum, frame):
raise TimeoutError(error_message)
def decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
signal(SIGALRM, _handle_timeout)
alarm(seconds)
try:
result = func(*args, **kwargs)
finally:
alarm(0)
return result
return wrapper
return decorator
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment