如何为生成序列切片的Python生成器函数添加正确类型标注

问题背景 #

我编写了如下chunk函数，需要为其添加精确的类型标注：

def chunk(data, chunk_size: int):
    yield from (data[i : i + chunk_size] for i in range(0, len(data), chunk_size))

需求如下：

• 支持bytes、bytearray或通用collections.abc.Sequence类型作为输入
• 返回类型为对应输入类型的生成器（如传入bytes时返回Generator[bytes]），而非宽泛的联合类型
• 避免使用typing.overload（代码冗余）

之前尝试的两种标注方式均触发pyright错误：

1. 标注为Sequence[T]返回Sequence[T]生成器：存在协变问题，传入bytes时无法匹配接收bytes的函数参数
2. 使用T: Sequence约束返回Generator[T]：pyright提示返回类型不兼容

解决方案 #

方案1：协变TypeVar + 类型转换 #

通过定义绑定到Sequence的协变TypeVar，并使用cast明确告知类型检查器切片返回的是输入类型本身，解决类型推断问题：

from collections.abc import Sequence, Generator
from typing import TypeVar, cast

T = TypeVar('T', bound=Sequence[Any], covariant=True)

def chunk(data: T, chunk_size: int) -> Generator[T, None, None]:
    for i in range(0, len(data), chunk_size):
        yield cast(T, data[i:i+chunk_size])

方案2：自定义Sliceable协议 #

定义一个协议来描述支持切片返回自身类型的序列，让类型检查器自动识别符合条件的类型（如bytes、bytearray、list），无需额外类型转换：

from collections.abc import Generator
from typing import TypeVar, Protocol

class SliceableSequence(Protocol):
    def __getitem__(self, key: slice) -> 'SliceableSequence': ...
    def __len__(self) -> int: ...

T = TypeVar('T', bound=SliceableSequence)

def chunk(data: T, chunk_size: int) -> Generator[T, None, None]:
    yield from (data[i:i+chunk_size] for i in range(0, len(data), chunk_size))

这两种方案都能满足需求：输入不同序列类型时返回对应类型的生成器，同时保持代码简洁。

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