super_flexible_decorator.py

from functools import wraps
from asyncio import iscoroutinefunction


def super_flexible_decorator(*deco_args, arg1=42, arg2=False):
    """
    A decorator that can be used with or without arguments, to decorate normal and async functions,
    and with any number of arguments.

    Examples:

    @super_flexible_decorator
    def foo():...

    @super_flexible_decorator(arg1=True)
    def foo():...

    @super_flexible_decorator
    async def foo():...

    @super_flexible_decorator(arg1=True)
    async def foo():...

    @super_flexible_decorator
    def foo(x, y):...

    @super_flexible_decorator(arg1=True)
    def foo(x, y):...

    @super_flexible_decorator
    async def foo(x, y):...

    @super_flexible_decorator(arg1=True)
    async def foo(x, y):...
    """
    def shared_decorator_logic(result):
        # the shared logic between the two (async and sync) possible decorated functions
        pass

    def decorator(func):
        if iscoroutinefunction(func):
            # decorating an async function so we must return one too, that await its results
            @wraps(func)
            async def new_func(*args, **kwargs):
                result = await func(*args, **kwargs)
                shared_decorator_logic(result)
                return result
        else:
            # decorating a normal function so we must return a normal decorated function
            @wraps(func)
            def new_func(*args, **kwargs):
                result = func(*args, **kwargs)
                shared_decorator_logic(result)
                return result

        return new_func

    if deco_args:
        # small hack to be able to decorate without arguments but without having to always put ()
        assert len(deco_args) == 1 and callable(deco_args[0]), "You must pass a callable to the decorator"
        return decorator(deco_args[0])
    else:
        return decorator