How to create a Python decorator that can wrap either coroutine or function?

Question

I am trying to make a decorator to wrap either coroutines or functions.

The first thing I tried was a simple duplicate code in wrappers:

Answer 1

May be you can find better way to do it, but, for example, you can just move your wrapping logic to some context manager to prevent code duplication:

import asyncio
import functools
import time
from contextlib import contextmanager


def duration(func):
    @contextmanager
    def wrapping_logic():
        start_ts = time.time()
        yield
        dur = time.time() - start_ts
        print('{} took {:.2} seconds'.format(func.__name__, dur))

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if not asyncio.iscoroutinefunction(func):
            with wrapping_logic():
                return func(*args, **kwargs)
        else:
            async def tmp():
                with wrapping_logic():
                    return (await func(*args, **kwargs))
            return tmp()
    return wrapper

Answer 2

For me the accepted answer by @mikhail-gerasimov was not working w/ async FastAPI methods (though it did work with normal and coroutine functions outside of FastAPI). However, I found this example on github that does work w/ fastapi methods. Adapted (slightly) below:

def duration(func):

    async def helper(func, *args, **kwargs):
        if asyncio.iscoroutinefunction(func):
            print(f"this function is a coroutine: {func.__name__}")
            return await func(*args, **kwargs)
        else:
            print(f"not a coroutine: {func.__name__}")
            return func(*args, **kwargs)

    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        start_ts = time.time()
        result = await helper(func, *args, **kwargs)
        dur = time.time() - start_ts
        print('{} took {:.2} seconds'.format(func.__name__, dur))

        return result

    return wrapper

Alternatively, if you want to keep the contextmanager, you can also do that:

def duration(func):
    """ decorator that can take either coroutine or normal function """
    @contextmanager
    def wrapping_logic():
        start_ts = time.time()
        yield
        dur = time.time() - start_ts
        print('{} took {:.2} seconds'.format(func.__name__, dur))

    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        if not asyncio.iscoroutinefunction(func):
            with wrapping_logic():
                return func(*args, **kwargs)
        else:
            with wrapping_logic():
                return (await func(*args, **kwargs))
    return wrapper

The difference between this and the accepted answer is not large. Mainly we just need to create an async wrapper and await the function if the function is a coroutine.

In my testing, this example code works in try/except blocks in your decorated function as well as with statements.

It's still not clear to me why the wrapper needs to be async for async FastAPI methods.