entry/exit decorator for async function

user_309433 0
Viewed 39

How may I write an equivalent decorator/wrapper for async functions, which will log when the async completes (returns or raises), without simply making it synchronous by awaiting it.

import functools

def trace_entry_exit(func):
    """non asyncio decorator"""
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print("entry")
        try:
            result = func(*args, **kwargs)
            print("returned")  ### if func is async, this will print ~immediately
        except Exception as ex:
            print("raised")
            raise
        return result ### without awaiting func(), this will simply return the awaitable object ~immediately

    return wrapper

I'd like returned or raised to only print when the function returns (or raises), but not block on the call in the decorator.

The same would apply to anything wrapping func, except the decorator use case is unique in that it's expected to be ~transparent to both func body, and whoever calls func(), as it is/was in the non-async case.

python-asyncio python-3.x python
1 Answers

Simply make your wrapper async. Then you can await the function, and it will work just like the wrapped function but with the print after it finishes:

import functools

def trace_entry_exit(func):
    """asyncio decorator"""
    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        print("entry")
        try:
            result = await func(*args, **kwargs)
            print("returned")
        except Exception as ex:
            print("raised")
            raise
        return result ### without awaiting func(), this will simply return the awaitable object ~immediately

    return wrapper

If you want it to work for either a coroutine or a synchronous function:

import functools
import asyncio

def trace_entry_exit(func):
    """asyncio or not decorator"""
    if asyncio.iscoroutinefunction(func):
        @functools.wraps(func)
        async def wrapper(*args, **kwargs):
            print("entry")
            try:
                result = await func(*args, **kwargs)
                print("returned")
            except Exception as ex:
                print("raised")
                raise
            return result
        return wrapper
    else:
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            print("entry")
            try:
                result = func(*args, **kwargs)
                print("returned")
            except Exception as ex:
                print("raised")
                raise
            return result
        return wrapper
user_198462190
Related