Recipe for universal decorator that accepts async function

[dimaqq]

I wonder if there's an elegant way to make a universal decorator that accepts both sync and async function... or classes and async functions for that matter.

[GrahamDumpleton]

Not sure what you mean. Can you give a pseudo example.

[dimaqq]

Here's what I got so far, for tornado style web request handlers, e.g. async def get(self): ...

@wrapt.decorator
def log_on_error(wrapped, instance, args, kwargs):
    """
    Special log on error, usage:

    Either:
    @trace_on_error
    class Foo(RequestHandler):
        async def get(self):
            if bad:
                raise HTTPError(401, "bad")

    Or:
    class Foo(RequestHandler):
        @trace_on_error
        async def get(self):
            if bad:
                raise HTTPError(401, "bad")
    """

    if instance is None and inspect.isclass(wrapped):
        # wrapping a class, called at instantiation time
        for verb in ("get", "post"):
            rv = wrapped(*args, **kwargs)
            # decorate get, post, options, etc.
            for verb in rv.SUPPORTED_METHODS:
                meth = getattr(rv, verb.lower())
                # expecting `async def get(self): ...`
                if inspect.iscoroutinefunction(meth):
                    setattr(rv, verb, log_on_error(meth))
        return rv
    elif inspect.ismethod(wrapped) and inspect.iscoroutinefunction(wrapped):
        # wrapping an instance method, called at call time
        coro = wrapped(*args, **kwargs)

        # additional custom wrapper for awaitables
        async def helper():
            try:
                await coro
            except HTTPError as e:
                logging.warning("This is a special extra log statement: %s", str(e))
                raise

        return helper()
    else:
        raise ValueError("Can only decorate classes or instance methods")

I find the async def helper kinda... inelegant and perhaps against the proxy philosophy of wrapt

[dimaqq]

The code above had some bugs, I've updated it.

[GrahamDumpleton]

I don't have a good simple answer, but whenever I look at how you are doing this, it just seems like the wrong way of going about it.

I would have had the check as to whether is normal function or coroutine done in a function called at the point the decorator is applied, rather than when the decorator is called. Much simplified, but something like:

@wrapt.decorator
def log_on_error_function(wrapped, instance, args, kwargs):
    ...

@wrapt.decorator
def log_on_error_coroutinewrapped, instance, args, kwargs):
    ...

def log_on_error(wrapped):
    if inspect.iscoroutinefunction(wrapped):
        return log_on_error_coroutinewrapped(wrapped)
    else:
        return log_on_error_function(wrapped)

[GrahamDumpleton]

I am going to close out this issue now, but first offer a more complete example for anyone else who comes across the issue.

import asyncio
import inspect
import time

import wrapt

def delay(duration):
    def wrapper(wrapped):
        @wrapt.decorator
        async def _async_delay(wrapped, instance, args, kwargs):
            print("ASYNC SLEEP", duration)
            await asyncio.sleep(duration)
            return await wrapped(*args, **kwargs)

        @wrapt.decorator
        def _sync_delay(wrapped, instance, args, kwargs):
            print("SYNC SLEEP", duration)
            time.sleep(duration)
            return wrapped(*args, **kwargs)

        if inspect.iscoroutinefunction(wrapped):
            return _async_delay(wrapped)
        else:
            return _sync_delay(wrapped)

    return wrapper

@delay(5)
def f():
    print("f()")

@delay(5)
async def g():
    print("g()")

async def main():
    return await g()

f()

asyncio.run(main())