使用协议参数对包装函数进行类型检查

Question

我正在编写一些代码，其中包含将函数从类型 T 提升到 Optional[T] 的包装函数，但是当我使用作为协议的类型时，出现了问题。

例如，我可以有这样的东西：

from typing import (
    TypeVar,
    Callable as Fn,
    Optional as Opt,
    Protocol,
    Any,
)
from functools import (
    wraps
)

_T = TypeVar('_T')

# Lifting a general function of type T x T -> T
def lift_op(f: Fn[[_T, _T], _T]) -> Fn[[Opt[_T], Opt[_T]], Opt[_T]]:
    """Lift op."""
    @wraps(f)
    def w(x: Opt[_T], y: Opt[_T]) -> Opt[_T]:
        return x if y is None else y if x is None else f(x, y)
    return w

将运算符 op: T x T -> T 提升到 Opt[T] x Opt[T] -> Opt[T]。（我已经将 Optional 和 Callable 缩短为 Opt 和 Fn 以获得更短的行，仅此而已。

大多数情况下，这似乎工作正常，但如果我有一个适用于受限于协议的泛型类型的函数，就会出现问题。

假设我有一个函数需要我的类型支持 <。然后我可以使用协议

# Protocol for types supporting <
class Ordered(Protocol):
    """Types that support < comparison."""

    def __lt__(self: Ord, other: Any) -> bool:
        """Determine if self is < other."""
        ...


Ord = TypeVar('Ord', bound=Ordered)

并定义一个min函数为

# Min for two optionals
def min_direct(x: Opt[Ord], y: Opt[Ord]) -> Opt[Ord]:
    return x if y is None else y if x is None else \
           y if y < x else x # on ties choose x

如果我用两个整数来调用它

# mypy accepts that ints are Opt[Ord]
min_direct(1, 2)  # No problem here

mypy 将接受 int 是一个 Opt[Ord].

但是如果我使用提升功能，它就会中断：

@lift_op
def lift_min(x: Ord, y: Ord) -> Ord:
    """Return min of x and y."""
    return y if y < x else x

# Now int is no longer an Opt[Ord]!
lift_min(1, 2)    # Type error for both args.

我收到错误

error: Argument 1 to "lift_min" has incompatible type "int"; expected "Optional[Ord]"
error: Argument 2 to "lift_min" has incompatible type "int"; expected "Optional[Ord]"

显然 int 在这种情况下不是 Opt[Ord]。

如果我专门为int

写一个min函数就好了

# Lifting a function with a concrete type (int) instead
# of a protocol
@lift_op
def imin(x: int, y: int) -> int:
    """Hardwire type to int."""
    return x if x <= y else y

# Again int is Opt[Ord]
imin(1, 2)  # All is fine here...

或者如果我明确指定包装函数的类型：

# Try setting the type of the lifted function explicitly    
def lift_min_(x: Ord, y: Ord) -> Ord:
    """Return min of x and y."""
    return y if y < x else x

f: Fn[[Opt[Ord],Opt[Ord]], Opt[Ord]] = lift_op(lift_min_)
f(1, 2) # No problem here

我怀疑 lift_op 包装器的 return 类型与上面 f 的类型注释 Fn[[Opt[Ord],Opt[Ord]],Opt[Ord]] 不同，但我不是确定以什么方式。这不是 wraps() 调用，这没有什么区别。但也许 Ord 类型以某种方式被绑定，然后被不同地解释？

我不知道，我也不知道怎么弄明白。我需要做什么才能使包装函数工作，以便它接受 int，比如说，满足协议 Opt[Ord]?

如果你想要上下文中的代码，here is a playground

Answer 1

不幸的是，这看起来像 mypy 错误。您的 Ord 类型 var 在装饰器级别得到解析，因为 T 出现在装饰器 def 的左侧和右侧。它是部分正确的（您确实想确认 T 在原始函数和转换后的函数中是相同的），但是使该变量未绑定。这就是为什么您尝试进行作业的原因：当您这样做时

f: Fn[[Opt[Ord], Opt[Ord]], Opt[Ord]] = lift_op(lift_min_)

你让 Ord 再次绑定。我认为你应该重新检查 mypy 问题跟踪器并提交这个错误，如果它之前没有完成（我找不到）。

要重现此错误，您甚至可以不受限制地使用简单类型变量：

from typing import TypeVar, Callable as Fn, Optional as Opt

_T = TypeVar('_T')

def allow_none(f: Fn[[_T], _T]) -> Fn[[Opt[_T]], Opt[_T]]:
    return f  # type: ignore

@allow_none
def simple(x: _T) -> _T:
    return x

reveal_type(simple)  # N: Revealed type is "def (Union[_T`-1, None]) -> Union[_T`-1, None]"
simple(1)  # E: Argument 1 to "simple" has incompatible type "int"; expected "Optional[_T]"

现在几乎任何类型（除了 Any 和特定类型的变量）都不能作为 simple 参数传递，因为它不是 _T。绑定变量表示为 def [_T] (Union[_T`-2, None]) -> Union[_T`-2, None].

这里是 related mypy issue，但它并没有完全涵盖您的情况，所以单独报告会很好（这可能会导致维护者提高优先级和更快的修复）。这个错误也会在 pyre 中重现，所以我可能误解了什么 - 但它看起来真的很奇怪，应该在 Callable 中绑定的类型变量不再绑定。

我尝试了一些涉及 Generic 和 Protocol 的解决方案，但其中 none 似乎有效：函数定义中的 type var binding is strict enough and callable with type variable from Generic 表现得非常奇怪，在 Union 上下文中解析为 <nothing> （不要因为小写 class 命名而杀了我，它与 classmethod 太相似了或者property 并故意保留小写）：

class _F(Protocol[_T]):
    def __call__(self, __x: _T) -> _T:
        ...

# Fails
class allow_none(Generic[_T]):
    def __call__(self, f: _F[_T], /) -> _F[_T | None]:
        return f  # type: ignore

reveal_type(allow_none[Ord]().__call__)  # N: Revealed type is "def (__main__._F[Ord?]) -> __main__._F[None]"

# But this works for some reason
class allow_none(Generic[_T]):
    def __call__(self, f: _F[_T], /) -> _F[_T]:
        return f  # type: ignore

reveal_type(allow_none[Ord]().__call__)  # N: Revealed type is "def (__main__._F[Ord?]) -> __main__._F[Ord?]"

估计也是bug，我会尽快反馈的。

不幸的是，对我来说，您最初尝试使用两步定义看起来是最好的解决方法。

我还建议看一下 Nikita Sobolev 的 returns library - 这个 link 指向处理您的用例的 Maybe 容器。虽然我不太喜欢它，但有些人认为它是 is None 链检查的更好替代方案。它附带为此类情况（几乎）免费打字。

使用协议参数对包装函数进行类型检查

Type checking wrapped functions with protocol arguments

python

mypy