Python中二维转置函数的正确类型注解方案

我们在代码中常使用zip(*something)实现二维转置，其中something是元组列表。例如：

>>> a = [('a', 1), ('b', 2), ('c', 3)]
>>> result = tuple(zip(*a))
>>> result  # 期望类型检查器识别为 tuple[tuple[str, ...], tuple[int, ...]]
(('a', 'b', 'c'), (1, 2, 3))

但在MyPy和Pyright中，tuple(zip(*a))的类型会被识别为tuple[Any]，完全丢失类型信息。现在需要实现一个transpose(iterable)函数，让类型检查器能正确识别返回类型，避免类型丢失。

利用Python的**可变泛型（Variadic Generics）**可以实现带精确类型注解的转置函数，让类型检查器准确推导返回类型。

实现代码 #

如果使用Python 3.11+，可以直接用标准库的类型；3.10及以下需要先安装typing-extensions库：

from typing import Iterable, Tuple
# Python 3.11+ 替换为 from typing import TypeVarTuple, Unpack
from typing_extensions import TypeVarTuple, Unpack

# 定义可变类型元组，用于捕获输入元组的各位置类型
Ts = TypeVarTuple('Ts')

def transpose(iterable: Iterable[Tuple[Unpack[Ts]]]) -> Tuple[Tuple[Ts[i], ...] for i in range(len(Ts))]:
    return tuple(zip(*iterable))

类型注解说明 #

• TypeVarTuple('Ts')：定义一组可变的类型变量，对应输入元组中每个位置的类型（比如输入元组是(str, int)，Ts就代表(str, int)这组类型）。
• Iterable[Tuple[Unpack[Ts]]]：表示输入是一个可迭代对象，其中每个元素都是包含Ts所有类型的元组。
• 返回类型Tuple[Tuple[Ts[i], ...] for i in range(len(Ts))]：表示转置后的结果是一个元组，每个元素是原输入对应位置所有元素组成的元组。比如输入是Iterable[tuple[str, int]]，返回类型就是tuple[tuple[str, ...], tuple[int, ...]]。

验证效果 #

用示例代码测试：

a = [('a', 1), ('b', 2), ('c', 3)]
result = transpose(a)
# 类型检查器会识别 result 的类型为 tuple[tuple[str, ...], tuple[int, ...]]

此时MyPy和Pyright都能正确推导类型，不会再出现tuple[Any]的情况。

注意事项 #

• 输入的所有元组必须长度一致，否则zip会截断到最短元组的长度，这和原生zip(*iterable)的行为一致，类型检查器目前无法处理长度不一致的输入。
• 如果需要返回迭代器而非元组，可以把返回类型改为Iterator[Tuple[Unpack[Tuple[Ts[i], ...] for i in range(len(Ts))]]]，函数内用yield from zip(*iterable)替代return tuple(...)。

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