added decorator_with_kwargs

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'
 ]

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