如何为基于元组长度返回不同类型的函数编写重载类型注解？

问题描述 #

我需要为一个根据输入元组长度返回不同类型的函数添加类型重载注解，函数原定义和语义如下：

def f(x: Tuple[int, ...]) -> Union[int, List[int]]:
    if len(x) == 1:
        return x[0]
    return list(x)

当输入元组长度为1时返回int，否则返回List[int]。我尝试了以下重载写法：

from typing import overload, List, Tuple

@overload
def f(x: Tuple[int]) -> int: ...

@overload
def f(x: Tuple[int, ...]) -> List[int]: ...

def f(x: Tuple[int, ...]) -> Union[int, List[int]]:
    if len(x) == 1:
        return x[0]
    return list(x)

但用mypy检查时出现错误：

script.py:4: error: Overloaded function signatures 1 and 2 overlap with incompatible return types
Found 1 error in 1 file (checked 1 source file)

问题在于第二个重载需要表示“长度不为1的任意元组”，但PEP 484本身不支持这类排他类型，请问正确的重载定义方式是什么？

解决方案 #

方法1：使用PEP 646的变长元组类型（推荐，Python 3.10+） #

利用PEP 646引入的TypeVarTuple和Unpack，可以精确定义“长度至少为2的元组”作为第二个重载的输入类型，彻底避免与第一个重载的Tuple[int]产生重叠：

from typing import overload, List, Tuple, TypeVarTuple, Unpack

Ts = TypeVarTuple('Ts')

@overload
def f(x: Tuple[int]) -> int: ...

@overload
def f(x: Tuple[int, int, Unpack[Ts]]) -> List[int]: ...

def f(x: Tuple[int, ...]) -> int | List[int]:
    if len(x) == 1:
        return x[0]
    return list(x)

这种写法能让类型检查器准确区分两种输入场景：

• 传入长度为1的元组（如(1,)）时，匹配第一个重载，推断返回int
• 传入长度≥2的元组（如(1,2)或(1,2,3)）时，匹配第二个重载，推断返回List[int]
  注意：需要使用支持PEP 646的类型检查器（如mypy 0.910+），且Python版本需≥3.10。

方法2：保留原有重载并抑制重叠错误（兼容旧版本） #

如果需要兼容Python 3.10以下版本，可以保留原来的重载结构，通过# type: ignore[misc]注释告诉mypy忽略重叠错误。因为类型检查器在处理调用时会优先匹配更具体的重载（即第一个Tuple[int]），实际类型推断依然有效：

from typing import overload, List, Tuple

@overload
def f(x: Tuple[int]) -> int: ...

@overload
def f(x: Tuple[int, ...]) -> List[int]: ...  # type: ignore[misc]

def f(x: Tuple[int, ...]) -> Union[int, List[int]]:
    if len(x) == 1:
        return x[0]
    return list(x)

这种写法的优势是兼容性强，但需要注意注释仅用于抑制误报的重叠错误，不会影响正常的类型推断逻辑。

错误原因解释 #

原写法报错是因为Tuple[int, ...]是变长元组类型，它包含了所有长度≥1的int元组，自然也覆盖了Tuple[int]（长度为1的元组）的情况。两个重载的输入类型存在重叠，但返回类型不同，类型检查器会判定为冲突。

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