同步和异步函数的类型提示装饰器

Question

如何输入用于同步和异步函数的提示装饰器？
我已经尝试过类似下面的方法，但是

mypy

会引发错误：

x/decorator.py:130: error: Incompatible types in "await" (actual type "Union[Awaitable[Any], R]", expected type "Awaitable[Any]")  [misc]
x/decorator.py:136: error: Incompatible return value type (got "Union[Awaitable[Any], R]", expected "R")  [return-value]

def log_execution_time(foo: Callable[P, AR | R]) -> Callable[P, AR | R]:
    module: Any = inspect.getmodule(foo)
    module_spec: Any = module.__spec__ if module else None
    module_name: str = module_spec.name if module_spec else foo.__module__  # noqa

    @contextmanager
    def log_timing():
        start = time()
        try:
            yield
        finally:
            exec_time_ms = (time() - start) * 1000
            STATS_CLIENT.timing(
                metric_key.FUNCTION_TIMING.format(module_name, foo.__name__),
                exec_time_ms,
            )

    async def async_inner(*args: P.args, **kwargs: P.kwargs) -> R:
        with log_timing():
            result = await foo(*args, **kwargs)  <- error
        return result

    def sync_inner(*args: P.args, **kwargs: P.kwargs) -> R:
        with log_timing():
            result = foo(*args, **kwargs)
        return result  <- error

    if inspect.iscoroutinefunction(foo):
        return wraps(foo)(async_inner)
    return wraps(foo)(sync_inner)

我知道有这样的技巧：

    if inspect.iscoroutinefunction(foo):
        async_inner: foo  # type: ignore[no-redef, valid-type]
        return wraps(foo)(async_inner)
    sync_inner: foo  # type: ignore[no-redef, valid-type]
    return wraps(foo)(sync_inner)

但我希望有一种方法可以正确输入提示。

我使用的是 python 3.10.10。

PS。我忘了说 PyCharm 挑选它并建议正确的类型很重要。

Answer 1

假设您想编写一个函数装饰器，在实际函数调用之前和/或之后执行some操作。我们将其称为周围环境

my_context

。如果您希望装饰器适用于异步函数和常规函数，则需要在其中容纳这两种类型。

我们如何正确地对其进行注释以确保类型安全性和一致性，同时还保留被包装函数的所有可能的类型信息？

这是我能想到的最干净的解决方案：

from collections.abc import Awaitable, Callable
from functools import wraps
from inspect import iscoroutinefunction
from typing import ParamSpec, TypeVar, cast, overload


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


@overload
def decorator(func: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]: ...


@overload
def decorator(func: Callable[P, R]) -> Callable[P, R]: ...


def decorator(func: Callable[P, R]) -> Callable[P, R] | Callable[P, Awaitable[R]]:
    if iscoroutinefunction(func):
        async def async_inner(*args: P.args, **kwargs: P.kwargs) -> R:
            with my_context():
                result = await cast(Awaitable[R], func(*args, **kwargs))
            return result
        return wraps(func)(async_inner)

    def sync_inner(*args: P.args, **kwargs: P.kwargs) -> R:
        with my_context():
            result = func(*args, **kwargs)
        return result
    return wraps(func)(sync_inner)

像这样应用它：

@decorator
def foo(x: str) -> str: ...


@decorator
async def bar(y: int) -> int: ...

这通过了

mypy --strict

，并在装饰后显示了

foo

和

bar

的类型，显示了我们所期望的，即分别为

def (x: builtins.str) -> builtins.str

和

def (y: builtins.int) -> typing.Awaitable[builtins.int]

。

详情

问题是，无论我们在内部包装器的outside应用什么类型的防护，它们都不会转移到包装器的inside。这是 mypy 的

长期存在的问题

，处理起来并不简单。

在我们的场景中，这意味着在装饰器中完成但内部包装器的outside

的任何

inspect.iscoroutinefunction检查只会将类型缩小到装饰器范围内的awaitable，但在包装器内部会被忽略。 （原因是即使在包装器定义之后，也可以对保存函数引用的非局部变量进行赋值。有关详细信息/示例，请参阅问题线程。）

在我看来，

typing.cast

是最直接的解决方法。

typing.overload

是一种在调用方处理协程和普通函数之间区别的方法。但我很好奇，想看看其他人能想出什么。