python - 我可以使用另一个函数的类型信息作为 Python 中的返回类型吗？

Question

在下面的示例中，如何正确注释 sum_two 函数的返回类型？

from typing import Any, TypeVar

T = TypeVar('T')
S = TypeVar('S')

def sum_two(first: T, second: S):
    return first + second

假设 __add__ 运算符已为将传递给此函数的所有可能参数进行了正确注释，是否有某种方法将返回类型表示为调用 __add__ 的返回类型code> 位于 T 和 S 类型的对象上？

我想避免使用打字的overload装饰器来识别所有可能的情况，因为可能有几十种情况。

Answer 1

理论上，您可以通过将 first 设为通用协议(protocol)来实现他的部分功能，该协议(protocol)可以让您“捕获”__add__ 的返回类型。例如:

# If you are using Python 3.7 or earlier, you'll need to pip-install
# the typing_extensions module and import Protocol from there.
from typing import TypeVar, Protocol, Generic

TOther = TypeVar('TOther', contravariant=True)
TSum = TypeVar('TSum', covariant=True)

class SupportsAdd(Protocol, Generic[TOther, TSum]):
    def __add__(self, other: TOther) -> TSum: ...

然后，您可以执行以下操作:

S = TypeVar('S')
R = TypeVar('R')

# Due to how we defined the protocol, R will correspond to the
# return type of `__add__`.
def sum_two(first: SupportsAdd[S, R], second: S) -> R:
    return first + second

# Type checks
reveal_type(sum_two("foo", "bar"))  # Revealed type is str
reveal_type(sum_two(1, 2))          # Revealed type is int
reveal_type(sum_two(1.0, 2))        # Revealed type is float

# Does not type check, since float's __radd__ is ignored
sum_two(1, 2.0)

class Custom:
    def __add__(self, x: int) -> int:
        return x

# Type checks
reveal_type(sum_two(Custom(), 3))  # Revealed type is int

# Does not type check
reveal_type(sum_two(Custom(), "bad"))

但是，这种方法确实有一些限制:

1. 它不处理“first”中没有匹配的 __add__ 但在“second”中具有匹配的 __radd__ 的情况。
2. 如果修改 Custom，则可能会得到一些奇怪的结果，因此 __add__ 是一个重载。我认为至少 mypy 目前有一个错误，它不知道如何正确处理涉及子类型和重载的复杂情况。