无代码重复的Python异步和同步函数装饰器

Question

我看到过至少两个例子，这些例子展示了如何制作一个可以同时处理普通函数def sync_func()和异步函数async def async_function()的装饰器。

不过，这些例子基本上都是重复了装饰器的代码，像这样：

import asyncio
import typing as t
import functools

R = t.TypeVar("R")
P = t.ParamSpec("P")


def decorator(func: t.Callable[P, R]) -> t.Callable[P, R]:
    if asyncio.iscoroutinefunction(func):
        @functools.wraps(func)
        async def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
            # … do something before
            result = await func(*args, **kwargs)
            # … do something after
            return result
    else:
        @functools.wraps(func)
        def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
            # … do something before [code duplication!]
            result = func(*args, **kwargs)
            # … do something after [code duplication!]
            return result
        
    return decorated

有没有可能创建一个装饰器，不在decorated里面重复代码呢？

Answer 1

你可以利用 Couroutine.send 的可迭代特性来运行一个 async def 函数，而不需要使用事件循环。

这个想法来源于 cpython上的一个评论。

import asyncio
import typing as t
import functools
import time

R = t.TypeVar("R")
P = t.ParamSpec("P")


def is_coroutine(
    func: t.Callable[P, t.Any]
) -> t.TypeGuard[t.Callable[P, t.Coroutine[t.Any, t.Any, t.Any]]]:
    """Return true if the given func is a coroutine"""
    # partial call needs to be removed
    while isinstance(func, functools.partial):
        func = func.func
    if asyncio.iscoroutinefunction(func):
        return True
    if hasattr(func, "__call__"):
        # handle classes with `async def __call__(…)` as well
        func = func.__call__
        return asyncio.iscoroutinefunction(func)
    return False



def decorator(func: t.Callable[P, R]) -> t.Callable[P, R]:
    is_sync = not is_coroutine(func)

    @functools.wraps(func)
    async def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
        print(f"{is_sync=}")
        print("Before")
        result = func(*args, **kwargs)
        if not is_sync:
            result = await t.cast(t.Awaitable[R], result)
        print(f"After with {result=}")
        return result

    if is_sync:
        @functools.wraps(func)
        def decorated_sync(*args: P.args, **kwargs: P.kwargs) -> R:
            """
            Execute coroutine by yield from all `await` calls without using event loop
            """
            try:
                decorated(*args, **kwargs).send(None)
            except StopIteration as result:
                return t.cast(R, result.value)
            raise RuntimeError("Never") # just to make mypy happy

        return decorated_sync

    return decorated  # type: ignore[return-value]

⚠️ 注意

尽管 async def decorated(…) 看起来像一个普通的 async 函数，并且你可以 await 其他的 async def 函数，但如果 is_sync 为 True，它们无法访问事件循环（包括 asyncio.* 里的内容或使用它的函数）。

请注意这一点！

最好在装饰器中只使用 同步 代码，或者在 if not is_sync: 之后再使用异步代码。

用法

它的工作方式和你预期的一样。

@decorator
def sync_call() -> str:
    time.sleep(0.1)
    return "sync"


@decorator
async def async_call() -> str:
    await asyncio.sleep(0.1)
    return "async"


async def run() -> None:
    res = sync_call()
    print(res)

    res = await async_call()
    print(res)


asyncio.run(run())

这导致了

is_sync=True
Before
After with result='sync'
sync
is_sync=False
Before
After with result='async'
async

可以在 MyPy Play 上查看