在Python中为具有条件输出类型的函数装饰器输入类型

Question

我有一组函数，它们都接受一个叫做 value 的参数，还有一些其他的命名参数。

我有一个装饰器：lazy。通常情况下，被装饰的函数会正常返回结果，但如果 value 是 None，它就会返回一个部分函数。

我该如何给这个装饰器添加类型提示，因为它的输出取决于输入的值呢？

from functools import partial

def lazy(func):
    def wrapper(value=None, **kwargs):
        if value is not None:
            return func(value=value, **kwargs)
        else:
            return partial(func, **kwargs)
    return wrapper

@lazy
def test_multiply(*, value: float, multiplier: float) -> float:
    return value * multiplier

@lazy
def test_format(*, value: float, fmt: str) -> str:
    return fmt % value

print('test_multiply 5*2:', test_multiply(value=5, multiplier=2))
print('test_format 7.777 as .2f:', test_format(value=7.777, fmt='%.2f'))

func_mult_11 = test_multiply(multiplier=11)  # returns a partial function
print('Type of func_mult_11:', type(func_mult_11))
print('func_mult_11 5*11:', func_mult_11(value=5))

我正在使用 mypy，并且通过 mypy 的扩展功能已经解决了大部分问题，但在 wrapper 中还没有搞定 value 的类型提示：

from typing import Callable, TypeVar, ParamSpec, Any, Optional
from mypy_extensions import DefaultNamedArg, KwArg

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

def lazy(func: Callable[P, R]) -> Callable[[DefaultNamedArg(float, 'value'), KwArg(Any)], Any]:
    def wrapper(value = None, **kwargs: P.kwargs) -> R | partial[R]:
        if value is not None:
            return func(value=value, **kwargs)
        else:
            return partial(func, **kwargs)
    return wrapper

我该如何给 value 添加类型提示？更好的是，我能在不使用 mypy 扩展的情况下做到这一点吗？

Answer 1

这里有两个可能的选择。第一个是“更正式的正确”，但太宽松了，依赖于 partial 提示的方法：

from __future__ import annotations

from functools import partial
from typing import Callable, TypeVar, ParamSpec, Any, Optional, Protocol, overload, Concatenate

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

class YourCallable(Protocol[P, R]):
    @overload
    def __call__(self, value: float, *args: P.args, **kwargs: P.kwargs) -> R: ...
    @overload
    def __call__(self, value: None = None, *args: P.args, **kwargs: P.kwargs) -> partial[R]: ...

def lazy(func: Callable[Concatenate[float, P], R]) -> YourCallable[P, R]:
    def wrapper(value: float | None = None, *args: P.args, **kwargs: P.kwargs) -> R | partial[R]:
        if value is not None:
            return func(value, *args, **kwargs)
        else:
            if args:
                raise ValueError("Lazy call must provide keyword arguments only")
            return partial(func, **kwargs)
    return wrapper  # type: ignore[return-value]

@lazy
def test_multiply(value: float, *, multiplier: float) -> float:
    return value * multiplier

@lazy
def test_format(value: float, *, fmt: str) -> str:
    return fmt % value

print('test_multiply 5*2:', test_multiply(value=5, multiplier=2))
print('test_format 7.777 as .2f:', test_format(value=7.777, fmt='%.2f'))

func_mult_11 = test_multiply(multiplier=11)  # returns a partial function
print('Type of func_mult_11:', type(func_mult_11))
print('func_mult_11 5*11:', func_mult_11(value=5))
func_mult_11(value=5, multiplier=5)  # OK
func_mult_11(value='a')  # False negative: we want this to fail

最后两个调用展示了这种方法的优缺点。partial 可以接受任何输入参数，所以并不够安全。如果你想覆盖最初提供给懒调用的参数，这可能是最好的解决方案。

注意，我稍微改变了输入调用的签名：如果不这样做，你将无法使用 Concatenate。另外，KwArg、DefaultNamedArg 等都已经被弃用，建议使用协议。你不能只用 paramspec 和 kwargs，args 也必须存在。如果你信任你的类型检查器，使用仅有 kwargs 的调用是可以的，所有未命名的调用将在类型检查阶段被拒绝。

不过，如果你不想覆盖传递给初始调用的默认参数，我还有另一个完全安全的替代方案，但如果你尝试这样做，它会产生误报。

from __future__ import annotations

from functools import partial
from typing import Callable, TypeVar, ParamSpec, Any, Optional, Protocol, overload, Concatenate

_R_co = TypeVar("_R_co", covariant=True)
R = TypeVar("R")
P = ParamSpec("P")

class ValueOnlyCallable(Protocol[_R_co]):
    def __call__(self, value: float) -> _R_co: ...
    
class YourCallableTooStrict(Protocol[P, _R_co]):
    @overload
    def __call__(self, value: float, *args: P.args, **kwargs: P.kwargs) -> _R_co: ...
    @overload
    def __call__(self, value: None = None, *args: P.args, **kwargs: P.kwargs) -> ValueOnlyCallable[_R_co]: ...


def lazy_strict(func: Callable[Concatenate[float, P], R]) -> YourCallableTooStrict[P, R]:
    def wrapper(value: float | None = None, *args: P.args, **kwargs: P.kwargs) -> R | partial[R]:
        if value is not None:
            return func(value, *args, **kwargs)
        else:
            if args:
                raise ValueError("Lazy call must provide keyword arguments only")
            return partial(func, **kwargs)
    return wrapper  # type: ignore[return-value]

@lazy_strict
def test_multiply_strict(value: float, *, multiplier: float) -> float:
    return value * multiplier

@lazy_strict
def test_format_strict(value: float, *, fmt: str) -> str:
    return fmt % value

print('test_multiply 5*2:', test_multiply_strict(value=5, multiplier=2))
print('test_format 7.777 as .2f:', test_format_strict(value=7.777, fmt='%.2f'))

func_mult_11_strict = test_multiply_strict(multiplier=11)  # returns a partial function
print('Type of func_mult_11:', type(func_mult_11_strict))
print('func_mult_11 5*11:', func_mult_11_strict(value=5))
func_mult_11_strict(value=5, multiplier=5)  # False positive: OK at runtime, but not allowed by mypy. E: Unexpected keyword argument "multiplier" for "__call__" of "ValueOnlyCallable"  [call-arg]
func_mult_11_strict(value='a')  # Expected. E: Argument "value" to "__call__" of "ValueOnlyCallable" has incompatible type "str"; expected "float"  [arg-type]

如果你愿意，你也可以在 ValueOnlyCallable 定义中将 value 标记为仅限关键字参数，但我觉得对于只有一个参数的函数来说，这样做不太合理。

你可以在 这个平台上比较这两种方法。

如果你不想使用忽略注释，下面的详细选项应该可以工作。不过，我认为为了去掉一个忽略注释而增加的冗长并不值得——这要看你自己怎么决定。

def lazy_strict(func: Callable[Concatenate[float, P], R]) -> YourCallableTooStrict[P, R]:
    @overload
    def wrapper(value: float, *args: P.args, **kwargs: P.kwargs) -> R: ...
    @overload
    def wrapper(value: None = None, *args: P.args, **kwargs: P.kwargs) -> ValueOnlyCallable[R]: ...
    def wrapper(value: float | None = None, *args: P.args, **kwargs: P.kwargs) -> R | ValueOnlyCallable[R]:
        if value is not None:
            return func(value, *args, **kwargs)
        else:
            if args:
                raise ValueError("Lazy call must provide keyword arguments only")    
            return partial(func, **kwargs)
    return wrapper

这里还有 Pyright 的 这个平台，因为 mypy 没能在我最初的回答中找到错误，而 Pyright 找到了。