This commit is contained in:
@@ -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'
|
||||
]
|
||||
|
@@ -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,14 +92,16 @@ 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):
|
||||
@wraps(function)
|
||||
def result(*args, **kwargs):
|
||||
attempts = 0
|
||||
delay = initial_delay
|
||||
@@ -48,16 +118,16 @@ def retry(
|
||||
|
||||
return result
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
@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):
|
||||
@wraps(function)
|
||||
async def result(*args, **kwargs):
|
||||
attempts = 0
|
||||
delay = initial_delay
|
||||
@@ -74,10 +144,9 @@ def async_retry(
|
||||
|
||||
return result
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
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):
|
||||
@decorator_with_kwargs
|
||||
def tmpdir(f, argument_name='tmpdir'):
|
||||
@wraps(f)
|
||||
def result(*args, **kwargs):
|
||||
with TemporaryDirectory() as temp_dir:
|
||||
fun(*args, **kwargs, **{
|
||||
f(*args, **kwargs, **{
|
||||
argument_name: Path(temp_dir)
|
||||
})
|
||||
|
||||
return result
|
||||
|
||||
return wrapper
|
||||
|
Reference in New Issue
Block a user