python - Mypy:寻找一个普通函数的完美签名

Question

我正在尝试为以下函数(Python 3.6，mypy 0.521)提出完美的函数签名:

def avg(xs):
    it = iter(xs)
    try:
        s = next(it)
        i = 1
    except StopIteration:
        raise ValueError("Cannot average empty sequence")
    for x in it:
        s += x
        i += 1
    return s / i

这段代码的优点在于，它适用于 int、float、complex 的迭代，并产生正确的结果, 也适用于 datetime.timedelta。尝试添加签名时出现问题。我尝试了以下方法:

def avg(xs: t.Iterable[t.Any]) -> t.Any: ...

但是现在，调用者需要转换结果。

def avg(xs: t.Iterable[T]) -> T: ...

这会失败，因为 T 不支持加法和除法。

N = TypeVar("N", int, float, complex, datetime.timedelta)
def avg(xs: t.Iterable[N]) -> N: ...

失败，因为 int/int 是 float；使用 // 几乎所有其他内容都会给出错误的结果。也很糟糕，因为代码应该适用于其他类型，只要支持加法和除法。

N = TypeVar("N", float, complex, datetime.timedelta)
def avg(xs: t.Iterable[N]) -> N: ...

这几乎是完美的，但同样，如果有人后来决定向它扔四元数，mypy 会提示。

...然后我也在尝试使用 abc 和 typing.overload 但那让我无处可去。

mypy --strict 下最优雅的解决方案是什么？

Answer 1

因此，不幸的是，Python/PEP 484 中的数字系统目前有点乱。

我们在技术上有一个 "numeric tower"这应该代表一组 ABC，Python 中的所有“类数字”实体都应该遵守这些 ABC。

此外，Python 中的许多内置类型(例如 int、float、complex 和 timedelta) 不从类型化中的这些 ABC 继承——这意味着这些 ABC 基本上不可用(除非您定义了显式从这些 ABC 继承的自定义类型)。

使问题复杂化的是 numbers module is largely dynamically typed在 typeshed 中——大约一年前我开始尝试修复数字模块，我的记忆是当时的 mypy 还不够强大，无法准确地输入数字塔。

这种情况今天可能会得到解决，但这或多或少都没有实际意义，因为 mypy 最近实现了对协议(protocol)的实验性支持(例如结构类型)!事实证明，这正是我们解决您的问题并最终修复数字塔所需要的(一旦将协议(protocol)添加到 PEP 484 和打字模块)。

现在，您需要做的是:

1. 安装typing_extensions模块(python3 -m pip install typing_extensions)
2. 从 Github 安装最新版本的 mypy(运行 python3 -m pip install -U git+git://github.com/python/mypy.git)

然后我们可以像这样为“支持加或除”类型定义一个协议(protocol):

from datetime import timedelta

from typing import TypeVar, Iterable
from typing_extensions import Protocol

T = TypeVar('T')
S = TypeVar('S', covariant=True)

class SupportsAddAndDivide(Protocol[S]):
    def __add__(self: T, other: T) -> T: ...

    def __truediv__(self, other: int) -> S: ...

def avg(xs: Iterable[SupportsAddAndDivide[S]]) -> S:
    it = iter(xs)
    try:
        s = next(it)
        i = 1
    except StopIteration:
        raise ValueError("Cannot average empty sequence")
    for x in it:
        s += x
        i += 1
    return s / i

reveal_type(avg([1, 2, 3]))
reveal_type(avg([3.24, 4.22, 5.33]))
reveal_type(avg([3 + 2j, 3j]))
reveal_type(avg([timedelta(1), timedelta(2), timedelta(3)]))

根据需要，使用 mypy 运行它会产生以下输出:

test.py:27: error: Revealed type is 'builtins.float*'
test.py:28: error: Revealed type is 'builtins.float*'
test.py:29: error: Revealed type is 'builtins.complex*'
test.py:30: error: Revealed type is 'datetime.timedelta*'