我有一个函数，可以提取任意变量的信息并返回结果，这样就可以把它嵌入到代码中，而无需再改写另一行。定义如下：

from typing import TypeVar

T = TypeVar("T")

def extract_info(var: T) -> T:
    # extracting and logging infos on var

    return var

因此，你只需把任意变量 foo 替换为 extract_info(foo)，就不会打乱代码，因为它不会修改变量，且会原样返回。

现在我想把函数更新为支持任意数量的变量，至少一个。我希望新函数是这样的：

def extract_info(var, *other_vars):
    # extracting and logging infos on var and other_vars if they exist

    all_args = (var, *other_vars)
    return all_args[0] if len(all_args) == 1 else all_args

这样就可以很方便地使用：

# PREVIOUS CODE
x: Any = some_operation()
t: tuple[Any] = some_other_operation()

# NEW CODE
x: Any = extract_info(some_operation())
t: tuple[Any] = some_operation(*some_other_operation())

在类型注解方面，问题就来了。输出的大小取决于输入的规模，类型也一样，具体如下：

• 如果只提供一个参数，返回类型就是该参数的类型
• 如果提供了其他参数，返回类型将是参数类型的元组（可能是异构的）

首先，我有如下想法：

from typing import TypeVar, TypeVarTuple

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

def extract_info(var: T, *other_vars: tuple[*Ts]) -> tuple[T, *Ts]:
    #do the rest...

但这个方案只适用于传入多个参数的情况。

有没有办法考虑到可能的类型变化？我可以在返回注解的末尾添加 | T，这样可行吗？还是说我的设计太复杂，无法正确进行类型提示？

解决方案

你可以使用 typing.overload 根据参数的数量（或类型）来改变返回类型。同时，你对 *other_vars 的提示有点偏离；应该用 *Ts，而不是 tuple[*Ts]。

from typing import overload, reveal_type

@overload
def extract_info[T](var: T) -> T: ...
@overload
def extract_info[T, *Ts](var: T, *other_vars: *Ts) -> tuple[T, *Ts]: ...
def extract_info[T, *Ts](var: T, *other_vars: *Ts) -> T | tuple[T, *Ts]:
    ... # implementation

reveal_type(extract_info(0)) # int
reveal_type(extract_info(0, 0.0)) # tuple[int, float]
reveal_type(extract_info(0, 0.0, object())) # tuple[int, float, object]
reveal_type(extract_info()) # error

这在 mypy 和 Pyright 上通过类型检查。