From 436bd737fa10234fea6c90a892cd1a3e4ffdedef Mon Sep 17 00:00:00 2001 From: Walter Oggioni Date: Mon, 24 Jun 2024 21:18:08 +0800 Subject: [PATCH] added `decorator_with_kwargs` --- src/pwo/__init__.py | 6 +- src/pwo/private.py | 156 +++++++++++++++++++++++++++++++------------- 2 files changed, 116 insertions(+), 46 deletions(-) diff --git a/src/pwo/__init__.py b/src/pwo/__init__.py index 5c2e471..b0e772b 100644 --- a/src/pwo/__init__.py +++ b/src/pwo/__init__.py @@ -4,7 +4,8 @@ from .private import ( retry, async_test, ExceptionHandlerOutcome, - tmpdir + tmpdir, + decorator_with_kwargs ) from .maybe import Maybe @@ -15,5 +16,6 @@ __all__ = [ 'async_test', 'ExceptionHandlerOutcome', 'Maybe', - 'tmpdir' + 'tmpdir', + 'decorator_with_kwargs' ] diff --git a/src/pwo/private.py b/src/pwo/private.py index 2a18fdd..f6fc84a 100644 --- a/src/pwo/private.py +++ b/src/pwo/private.py @@ -3,8 +3,76 @@ from tempfile import TemporaryDirectory from pathlib import Path from enum import Enum, auto from typing import Callable +from inspect import signature from time import sleep from asyncio import sleep as async_sleep, Runner +from functools import wraps, partial + + +def decorator_with_kwargs(decorator: Callable) -> Callable: + """Decorator factory to give decorated decorators the skill to receive + optional keyword arguments. + + If a decorator "some_decorator" is decorated with this function: + + @decorator_with_kwargs + def some_decorator(decorated_function, kwarg1=1, kwarg2=2): + def wrapper(*decorated_function_args, **decorated_function_kwargs): + '''Modifies the behavior of decorated_function according + to the value of kwarg1 and kwarg2''' + ... + return wrapper + + It will be usable in the following ways: + + @some_decorator + def func(x): + ... + + @some_decorator() + def func(x): + ... + + @some_decorator(kwarg1=3) # or other combinations of kwargs + def func(x, y): + ... + + :param decorator: decorator to be given optional kwargs-handling skills + :type decorator: Callable + :raises TypeError: if the decorator does not receive a single Callable or + keyword arguments + :raises TypeError: if the signature of the decorated decorator does not + conform to: Callable, **keyword_arguments + :return: modified decorator + :rtype: Callable + """ + + @wraps(decorator) + def decorator_wrapper(*args, **kwargs): + if (len(kwargs) == 0) and (len(args) == 1) and callable(args[0]): + return decorator(args[0]) + if len(args) == 0: + return partial(decorator, **kwargs) + raise TypeError( + f'{decorator.__name__} expects either a single Callable ' + 'or keyword arguments' + ) + + signature_values = signature(decorator).parameters.values() + signature_args = [ + param.name for param in signature_values + if param.default == param.empty + ] + + if len(signature_args) != 1: + raise TypeError( + f'{decorator.__name__} signature should be of the form:\n' + f'{decorator.__name__}(function: typing.Callable, ' + 'kwarg_1=default_1, kwarg_2=default_2, ...) -> Callable' + ) + + return decorator_wrapper + _size_uoms = ('B', 'KiB', 'MiB', 'GiB', 'KiB') @@ -24,60 +92,61 @@ class ExceptionHandlerOutcome(Enum): CONTINUE = auto() +@decorator_with_kwargs def retry( + function, max_attempts: int = 3, multiplier: float = 2, initial_delay: float = 1.0, exception_handler: Callable[[Exception], ExceptionHandlerOutcome] = lambda _: ExceptionHandlerOutcome.CONTINUE ): - def wrapper(function): - def result(*args, **kwargs): - attempts = 0 - delay = initial_delay - while True: - attempts += 1 - try: - return function(*args, **kwargs) - except Exception as ex: - if attempts < max_attempts and exception_handler(ex) == ExceptionHandlerOutcome.CONTINUE: - sleep(delay) - delay *= multiplier - else: - raise ex + @wraps(function) + def result(*args, **kwargs): + attempts = 0 + delay = initial_delay + while True: + attempts += 1 + try: + return function(*args, **kwargs) + except Exception as ex: + if attempts < max_attempts and exception_handler(ex) == ExceptionHandlerOutcome.CONTINUE: + sleep(delay) + delay *= multiplier + else: + raise ex - return result - - return wrapper + return result +@decorator_with_kwargs def async_retry( + function, max_attempts: int = 3, multiplier: float = 2, initial_delay: float = 1.0, exception_handler=lambda _: ExceptionHandlerOutcome.CONTINUE ): - def wrapper(function): - async def result(*args, **kwargs): - attempts = 0 - delay = initial_delay - while True: - attempts += 1 - try: - return await function(*args, **kwargs) - except Exception as ex: - if attempts < max_attempts and exception_handler(ex) == ExceptionHandlerOutcome.CONTINUE: - await async_sleep(delay) - delay *= multiplier - else: - raise ex + @wraps(function) + async def result(*args, **kwargs): + attempts = 0 + delay = initial_delay + while True: + attempts += 1 + try: + return await function(*args, **kwargs) + except Exception as ex: + if attempts < max_attempts and exception_handler(ex) == ExceptionHandlerOutcome.CONTINUE: + await async_sleep(delay) + delay *= multiplier + else: + raise ex - return result - - return wrapper + return result def async_test(coro): + @wraps(coro) def wrapper(*args, **kwargs): with Runner() as runner: runner.run(coro(*args, **kwargs)) @@ -85,14 +154,13 @@ def async_test(coro): return wrapper -def tmpdir(argument_name='tmpdir'): - def wrapper(fun): - def result(*args, **kwargs): - with TemporaryDirectory() as temp_dir: - fun(*args, **kwargs, **{ - argument_name: Path(temp_dir) - }) +@decorator_with_kwargs +def tmpdir(f, argument_name='tmpdir'): + @wraps(f) + def result(*args, **kwargs): + with TemporaryDirectory() as temp_dir: + f(*args, **kwargs, **{ + argument_name: Path(temp_dir) + }) - return result - - return wrapper + return result