如何为fun函数添加类型注解，使mypy识别异构参数的精确元组返回类型？

问题背景 #

现有如下代码，使用单个绑定A的TypeVar定义可变参数函数fun：

from typing import TypeVar, reveal_type

class A: ...
class B(A): ...
class C(A): ...

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

def fun(*args: T) -> tuple[T, ...]: return args

# 单个参数
reveal_type(fun(A()))  # 推导为 tuple[A, ...]
reveal_type(fun(B()))  # 推导为 tuple[B, ...]
reveal_type(fun(C()))  # 推导为 tuple[C, ...]

# 多个同类型参数
reveal_type(fun(A(), A()))  # 推导为 tuple[A, ...]
reveal_type(fun(B(), B()))  # 推导为 tuple[B, ...]
reveal_type(fun(C(), C()))  # 推导为 tuple[C, ...]

# 多个不同类型参数
reveal_type(fun(A(), B()))  # 推导为 tuple[A, ...]
reveal_type(fun(B(), C()))  # 推导为 tuple[A, ...]
reveal_type(fun(C(), A()))  # 推导为 tuple[A, ...]

当前实现中，传入不同类型的参数时，mypy会统一将返回类型推导为tuple[A, ...]，无法精确对应输入参数的具体类型组合（比如tuple[A, B]、tuple[B, C]）。

解决方案 #

使用Python 3.11+引入的TypeVarTuple和Unpack特性，实现可变参数的精确类型推导：

from typing import TypeVarTuple, Unpack, reveal_type

class A: ...
class B(A): ...
class C(A): ...

# 定义一组绑定A的可变类型变量
Ts = TypeVarTuple('Ts', bound=A)

def fun(*args: Unpack[Ts]) -> tuple[Unpack[Ts]]:
    return args

# 单个参数
reveal_type(fun(A()))  # 推导为 tuple[A]
reveal_type(fun(B()))  # 推导为 tuple[B]
reveal_type(fun(C()))  # 推导为 tuple[C]

# 多个同类型参数
reveal_type(fun(A(), A()))  # 推导为 tuple[A, A]
reveal_type(fun(B(), B()))  # 推导为 tuple[B, B]
reveal_type(fun(C(), C()))  # 推导为 tuple[C, C]

# 多个不同类型参数
reveal_type(fun(A(), B()))  # 推导为 tuple[A, B]
reveal_type(fun(B(), C()))  # 推导为 tuple[B, C]
reveal_type(fun(C(), A()))  # 推导为 tuple[C, A]

说明 #

• TypeVarTuple('Ts', bound=A)：定义一组类型变量，每个变量都必须是A或其子类
• *args: Unpack[Ts]：将这组类型变量展开，作为函数的可变参数类型
• tuple[Unpack[Ts]]：将这组类型变量展开，作为返回元组的精确类型

此特性需要Python 3.11及以上版本，且mypy版本需支持（推荐0.991及以上）。

内容的提问来源于stack exchange，提问作者Vito De Tullio