decorator icon indicating copy to clipboard operation
decorator copied to clipboard

Create a decorator for both sync and async function

Open Doma1204 opened this issue 1 year ago • 1 comments

In Python without decorator package, if I want to create a custom decorator that supports both sync and async function, I can write the following:

import inspect

def custom_decorator(func):
    if inspect.iscoroutinefunction(func):
        # Support for async functions
        async def awrapper(*args, **kwargs):
            print("custom_decorator")
            return await func(*args, **kwargs)
        return awrapper
    else:
        # Support for sync functions
        def wrapper(*args, **kwargs):
            print("custom_decorator")
            return func(*args, **kwargs)
        return wrapper

However, with the decorator package, how can I write a custom decorator that supports both sync and async function?

I've tried this way: (minimal example)

from decorator import decorator
import asyncio

@decorator
def custom_decorator(func, *args, **kwargs):
    print("custom_decorator_start")
    return func(*args, **kwargs)

@custom_decorator
def sum(x, y):
    return x+y

@custom_decorator
async def asum(x, y):
    await asyncio.sleep(1)
    return x+y

In this way, I can call the async asum function with await asum(1, 2), but if I check whether asum is a coroutine function with inspect.iscoroutinefunction, it returns False.

I've also tried another way:

import inspect
from decorator import decorator
import asyncio

@decorator
async def custom_decorator(func, *args, **kwargs):
    print("custom_decorator_start")
    if inspect.iscoroutinefunction(func):
        return await func(*args, **kwargs)
    else:
        return func(*args, **kwargs)

@custom_decorator
def sum(x, y):
    return x+y

@custom_decorator
async def asum(x, y):
    await asyncio.sleep(1)
    return x+y

But this time, calling inspect.iscoroutinefunction(sum) and inspect.iscoroutinefunction(asum) will both return True.

Is there any way I can write a custom decorator that supports both sync and async function with decorator package? If not, would you consider supporting this use case?

Doma1204 avatar Aug 08 '23 09:08 Doma1204

Python 3.12 finally added inspect.markcoroutinefunction to solve this.

async def foo():
    ...

inspect.markcoroutinefunction(foo)

Of course you could wrap decorator if you don't want to do this manually for every decorated async function.


Background

Python allows creating callables which are not coroutines (not defined with async) but still return awaitables (work with await).

So everyone who uses inspect.iscoroutinefunction to decide whether to use await is technically wrong (but Python itself has at least half the blame here - the need for something like inspect.markcoroutinefunction has been blatantly obvious for at least five years now, and arguably should've been obvious when first designing the async/await language addition back in 3.5).

inspect.iscoroutinefunction is a half check. If it returns True we know we should use await, but if it returns False we haven't actually disproven if we should use await.

So instead of

if inspect.iscoroutinefunction(f):
    result = await f(...)  # correct
else:
    result = f(...)  # could be wrong, f might return an awaitable

the universally correct solution is

result = f(...)
if inspect.isawaitable(result):
    result = await result

And if we're in a wrapper that needs to support both sync and async functions, you just return whatever f returned, which makes you transparent w.r.t. to whether an await is needed on your return value:

return f(...)

I infer decorator is using that last method, which is technically the most correct thing to do, but sadly it just pushes this obscure edge-case further up the call-stack, making it everyone else's problem.

So decorator could be written to progressively enhance: try importing inspect.markcoroutinefunction - if it's not available just don't use it, but if it is, check the wrapped callable with inspect.iscoroutinefunction and mark the wrapper to match.

mentalisttraceur avatar Apr 07 '24 08:04 mentalisttraceur