如何为Python存根处理嵌套Protocol与泛型

我正在为Stream对象定义一个基于输入、输出双泛型的Protocol，已实现map和sliding_window等方法，这些方法会返回新流，新流的输入为原流的输出，输出为其他类型。现有代码如下：

from __future__ import annotations

from typing import (
    Callable,
    Protocol,
    TypeVar,
    Tuple,
)

U = TypeVar("U", contravariant=True)
V = TypeVar("V")
T = TypeVar("T")


class StreamInterface(Protocol[U, V]):
    def map(
        self: StreamInterface[U, V],
        func: Callable[[V], T],
    ) -> StreamInterface[V, T]:
        ...

    def sliding_window(
        self: StreamInterface[U, V],
        n: int,
    ) -> StreamInterface[V, Tuple[V, ...]]:
        ...

现在需要定义flatten操作：仅当原流的输出类型是可迭代对象时才能调用该方法，调用后返回的新流会将原流输出的可迭代对象展开，输出其中的元素。我尝试用以下方式定义，但mypy无法正确跟踪类型：

from typing import Iterable

K = TypeVar("K")
IterableK = Iterable[K]
M = TypeVar("M", bound=IterableK)

# 在StreamInterface中添加：
def flatten(
    self: StreamInterface[U, M],
) -> StreamInterface[M, K]:
    ...

测试代码中，mypy会提示flattened需要类型注解：

def func1(input: int) -> float:
    return input + 1.2

def func2(input: float) -> str:
    return str(input)

source: StreamInterface[int, int] = Stream()
slider = source.map(func1).sliding_window(3)
flattened = slider.flatten()  # mypy报错：需要类型提示
flattened.map(func2).sink(print)

我猜测问题出在K、M与原泛型参数U、V的关联不足，但直接用self: StreamInterface[U, Iterable[V]]也无法被正确处理，想知道标注flatten方法类型的最佳方式。

解决方案 #

核心是让类型检查器明确关联原流输出的可迭代对象和展开后的元素类型，可以通过在flatten方法中引入专门的类型变量来约束原流的输出类型：

from __future__ import annotations

from typing import (
    Callable,
    Protocol,
    TypeVar,
    Tuple,
    Iterable,
)

U = TypeVar("U", contravariant=True)
V = TypeVar("V")
T = TypeVar("T")
# 用于表示flatten后输出的元素类型
K = TypeVar("K")


class StreamInterface(Protocol[U, V]):
    def map(
        self: StreamInterface[U, V],
        func: Callable[[V], T],
    ) -> StreamInterface[V, T]:
        ...

    def sliding_window(
        self: StreamInterface[U, V],
        n: int,
    ) -> StreamInterface[V, Tuple[V, ...]]:
        ...

    def flatten(
        # 约束当前流的输出V必须是K类型的可迭代对象
        self: StreamInterface[U, Iterable[K]],
    ) -> StreamInterface[Iterable[K], K]:
        ...

原理说明

• 直接在self的类型注解中约束V为Iterable[K]，让mypy能明确推断出：原流输出的是包含K类型元素的可迭代对象，flatten后新流的输出就是K类型。
• 这种方式不需要额外定义M类型变量，直接通过Iterable[K]建立原流输出与展开后元素的类型关联，mypy可以正确跟踪类型推导链。

修改后，测试代码中的flattened会被自动推断为StreamInterface[Tuple[float, ...], float]，mypy不再报错，后续调用map(func2)也能正常进行类型检查。

内容的提问来源于stack exchange，提问作者DIN14970