CODEWITHSUNDEEP
pythonpython-3.xasynchronous
Amin Etesamian
Amin Etesamian

Reputation: 3699

Using a coroutine as decorator

in this scenario:

async def foo(f):
    async def wrapper(*args, **kwargs):
        return f(*args, **kwargs)
    return wrapper

@foo
async def boo(*args, **kwargs):
    pass

is the call to foo as a decorator for boo decorator an async call?

--First Edit: Also how does one handle calling chain of coroutines as decorators?

Upvotes: 68

Views: 73584

Answers (5)

Mike Smith
Mike Smith

Reputation: 456

async def foo(f):
  async def wrapper(*args, **kwargs):
    # wrapper pre-function stuff
    result = await f(*args, **kwargs) # key is to await function's result
    # wrapper post-function stuff
    return result 
  wrapper.__name__ = f.__name__ # for some reason, async wrappers don't do this
  # do it to avoid an error if you use the wrapper on multiple functions
  return wrapper

The two key changes are to await the function you are wrapping, as it is an async function, and to change the name of the wrapper function so your program isn't trying to name multiple functions the same thing.

Upvotes: -1

Kiran Jonnalagadda
Kiran Jonnalagadda

Reputation: 2826

A decorator that works for both sync and async functions:

from inspect import iscoroutinefunction
from functools import wraps

def my_decorator(decorator_option=None):
    def decorator(f):
        def shared_logic(*args, **kwargs):
            # Do something with decorator_option, if relevant
            pass

        @wraps(f)
        def wrapper(*args, **kwargs):
            shared_logic(*args, **kwargs)
            return f(*args, **kwargs)

        @wraps(f)
        async def async_wrapper(*args, **kwargs):
            shared_logic(*args, **kwargs)
            return await f(*args, **kwargs)

        return async_wrapper if iscoroutinefunction(f) else wrapper
    return decorator

If the decorator doesn't take any args, the outer function can be removed.

Upvotes: 2

user11919393
user11919393

Reputation: 481

def foo(f):
    async def wrapper(*args, **kwargs):
        return await f(*args, **kwargs)
    return wrapper

@foo
async def boo(*args, **kwargs):
    pass

Your decorator needs to be a normal function and it will work fine.

When a decorator is evaluated python executes the method with the function as the argument.

@foo
async def boo():
    pass

Evaluates to:

__main__.boo = foo(boo)

If foo is an async function type(main.boo) will be a coroutine object, not a function object. But if foo is a regular synch function it will evaluate right away and main.boo will be the wrapper returned.

Upvotes: 48

Kevin Lyons
Kevin Lyons

Reputation: 43

Here is an alternate approach using the decorator library (i.e. pip install decorator first):

import asyncio

import decorator


@decorator.decorator
async def decorate_coro(coro, *args, **kwargs):
    try:
        res = await coro(*args, **kwargs)
    except Exception as e:
        print(e)
    else:
        print(res)


@decorate_coro
async def f():
    return 42


@decorate_coro
async def g():
    return 1 / 0


async def main():
    return await asyncio.gather(f(), g())

if __name__ == '__main__':
    asyncio.run(main())

Output:

42
division by zero

Upvotes: 3

Amin Etesamian
Amin Etesamian

Reputation: 3699

Thanks to @blacknght's comment, considering 

def foo():
    def wrapper(func):
        @functools.wraps(func)
        async def wrapped(*args):
             # Some fancy foo stuff
            return await func(*args)
        return wrapped
    return wrapper

and

def boo():
    def wrapper(func):
        @functools.wraps(func)
        async def wrapped(*args):
            # Some fancy boo stuff
            return await func(*args)
        return wrapped
    return wrapper

as two decorators, and

@foo()
@boo()
async def work(*args):
    pass

As the foo is wrapping the work coroutine, the key is to await the func(*arg) in both decorators.

Upvotes: 86

Related Questions