在 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 扩展来完成大部分工作，但还没有在

value

中进行

wrapper

打字：

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, **kwargs: P.kwargs) -> R | partial[R]:
        if value is not None:
            return func(value, **kwargs)
        else:
            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

和 company 均已弃用，取而代之的是协议。您不能仅将 paramspec 与 kwargs 一起使用，args 也必须存在。如果您信任您的类型检查器，则可以使用仅 kwarg 可调用对象，所有未命名的调用都将在类型检查阶段被拒绝。

但是，如果您不想覆盖传递给初始可调用对象的默认参数，我还有另一种选择可以分享，这是完全安全的，但如果您尝试这样做，则会发出误报。

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, **kwargs: P.kwargs) -> R | partial[R]:
        if value is not None:
            return func(value, **kwargs)
        else:
            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]

如果愿意，您还可以在

value

定义中标记

ValueOnlyCallable

kw-only，我只是认为对于只有一个参数的函数来说这是不合理的。

您可以在 playground 中比较这两种方法。