能否用Mypy对单个对象与解包元组进行函数重载？

你遇到的核心问题是：运行时很容易实现同时支持单个元素输入和解包多个元素的函数，但Mypy无法正确识别这种类型签名。你想保留foo(1, 2, ...)这种解包写法的简洁性，同时支持传入单个元素的场景（避免解包单个元组带来的额外索引操作），但当前的重载写法无法让类型检查器区分两种输入情况。

你的示例代码如下：

from typing import overload


@overload
def foo(a: int) -> int: # Impossible to distinguish inputs from overload below
    ...


@overload
def foo(*a: int) -> tuple[int, ...]:
    ...


def foo(*a: int | tuple[int, ...]) -> int | tuple[int, ...]:
    if len(a) == 1:
        return a[0]
    return a


assert foo(1) == 1 # This is the expected, but how would the type checker know?
assert foo(1, 2) == (1, 2) # This is obviously the correct signature

可行的解决方案 #

1. 用位置-only参数修复重载歧义

问题根源在于foo(1)会同时匹配两个重载签名，导致Mypy无法判断返回类型。通过将单参数重载标记为位置-only参数，可以让Mypy明确区分两种调用场景：

from typing import overload

@overload
def foo(a: int, /) -> int:
    ...

@overload
def foo(*a: int) -> tuple[int, ...]:
    ...

def foo(*a: int) -> int | tuple[int, ...]:
    if len(a) == 1:
        return a[0]
    return a

assert foo(1) == 1  # Mypy现在会识别为返回int
assert foo(1, 2) == (1, 2)  # 返回tuple[int, ...]

这里的/是Python的位置-only参数语法，运行时不影响调用方式，但会让Mypy将foo(1)优先匹配第一个重载，多参数调用匹配第二个重载。

2. 用类型守卫辅助类型推断

如果不想用位置-only参数，也可以通过显式的类型判断让Mypy理解返回值类型：

from typing import Union, Tuple

def foo(*a: int) -> Union[int, Tuple[int, ...]]:
    if len(a) == 1:
        return a[0]
    return a

# 调用后通过类型守卫让Mypy识别类型
result = foo(1)
if isinstance(result, int):
    print(result + 1)  # Mypy知道这里result是int
else:
    print(len(result)) # 这里result是tuple

这种写法的缺点是调用后需要显式判断类型，不如重载写法直观，但胜在简单直接。

3. 拆分函数（备选方案）

如果不想纠结类型标注的歧义，也可以拆分出两个明确的函数：比如foo_single(a: int) -> int和foo_multiple(*a: int) -> tuple[int, ...]，但这会失去你想要的调用简洁性，仅作为备选。

总结 #

完全避免解包不是唯一途径，通过位置-only参数配合重载就能让Mypy正确识别两种场景的类型，同时保留调用的简洁性。

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