Mypy 可以重载单个对象和解包元组吗？

Question

以下内容在运行时很容易实现，但在 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

完全避免拆包真的是唯一的方法吗？

Answer 1

*args

表示0个或更多位置参数，所以你需要更好的

@overload

签名：

• 如果仅传递一个参数，则返回该参数。
• 如果传递了至少 2 个参数（2 个具体 + 0 或更多），则返回此类参数的元组。

这些可以翻译为键入提示，如下所示：

_{（游乐场链接：mypy、Pyright、PEP 695 语法）}

from typing import overload, TypeVar, TypeVarTuple

T = TypeVar('T')
T2 =  TypeVar('T2')
Ts = TypeVarTuple('Ts')

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

@overload
def foo(a0: T, a1: T2, /, *rest: *Ts) -> tuple[T, T2, *Ts]:
    ...

def foo(a: T, /, *rest: *Ts) -> T | tuple[T, *Ts]:
    if len(rest) == 0:
        return a
    
    return (a, *rest)

reveal_type(foo(1))           # mypy & pyright => int
reveal_type(foo(1, 2))        # mypy & pyright => tuple[int, int]
reveal_type(foo(1, 2., '3'))  # mypy           => tuple[int, float, Literal['3']]
                              # pyright        => tuple[int, float, str]

foo()                         # error
foo(2, bar = 4)               # error