Python 3.12中如何为可变参数函数添加精确类型注解？

问题描述 #

我们需要给以下Python函数添加类型注解：

def f(*args: float) -> tuple[float, ...]:
    ...
    return tuple(args)

要求通过mypy或pylance实现两项检查：

• 返回元组的元素数量与传入的args完全一致
• 所有元素的类型均为float

已知使用TypeVarTuple可以约束返回元组的元素数量与参数匹配，但无法限制元素类型。除了为1参数、2参数等场景编写大量重载外，Python 3.12中是否有更简洁的实现方式？

解决方案 #

可以结合TypeVarTuple、TypeVar和Unpack实现需求，无需编写大量重载。具体代码如下：

from typing import TypeVar, TypeVarTuple, Unpack

# 定义绑定到float的类型变量，限定元素类型范围
T = TypeVar('T', bound=float)
# 定义类型元组，约束所有元素必须符合T的类型要求
FloatArgs = TypeVarTuple('FloatArgs', bound=tuple[T, ...])

def f(*args: Unpack[FloatArgs]) -> FloatArgs:
    ...
    return tuple(args)

实现说明 #

• TypeVar('T', bound=float) 限定了类型元组中每个元素必须是float或其子类型
• TypeVarTuple('FloatArgs', bound=tuple[T, ...]) 定义了一个类型元组FloatArgs，其所有元素都满足T的类型约束
• *args: Unpack[FloatArgs] 表示可变参数args是FloatArgs类型元组展开后的元素，确保输入参数的数量和类型合法
• 返回类型FloatArgs直接绑定到输入的类型元组，保证返回元组的元素数量、类型与输入完全匹配

检查效果 #

使用mypy或pylance检查时：

• 若函数内部返回的元组元素数量与输入参数数量不一致，会触发类型错误
• 若传入非float类型的参数，或返回元组包含非float元素，也会触发类型错误

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