Last active
August 17, 2024 22:38
-
-
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.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
"""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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
"""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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
"""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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
"""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() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
"""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() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
"""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