Python可变参数函数如何实现依赖入参类型的返回值类型提示

想要实现「入参可变长度、返回值元组每个位置类型和入参一一对应」的标注，目前有两种成熟可行的实现方式：

方案一：使用TypeVarTuple（推荐，支持任意长度参数） #

这个特性是PEP 646正式引入的，专门解决可变泛型参数的类型映射问题，Python 3.11+可以直接从标准库typing导入相关对象，更低版本可以通过typing_extensions兼容包导入，目前mypy、pyright等所有主流类型检查器都已经完整支持。

标注后的完整代码如下：

from typing import TypeVar, TypeVarTuple, Unpack, Any
# Python <3.11 请替换为下面的导入语句
# from typing_extensions import TypeVarTuple, Unpack
from collections.abc import Mapping, Hashable
from functools import reduce

K = TypeVar('K', bound=Hashable)
# 声明可变长度类型占位序列，对应每个入参Mapping的value类型
Vs = TypeVarTuple('Vs')

def keyset(m: Mapping[K, Any]) -> set[K]:
    return set(m.keys())

def inner_join(*ms: Unpack[tuple[Mapping[K, v] for v in Vs]]) -> dict[K, tuple[Unpack[Vs]]]:
    keys = reduce(set.intersection, map(keyset, ms)) if ms else set()
    return {k: tuple(m[k] for m in ms) for k in keys}

这个标注的效果和固定两参数版本完全一致：类型检查器会按位置匹配每个入参Mapping的value类型，自动推导返回值元组对应位置的类型。比如传入inner_join({"id": 1}, {"id": "admin"}, {"id": True})时，会自动推导返回值类型为dict[str, tuple[int, str, bool]]，没有精度损失。

方案二：使用@overload重载固定参数个数（兼容旧环境） #

如果需要兼容不支持PEP 646的旧Python版本或者旧类型检查器，可以用重载的方式，提前定义好常用参数个数的精确类型，超出个数的场景用兜底的弱精度标注即可。
示例代码如下：

from typing import TypeVar, overload, Any
from collections.abc import Mapping, Hashable
from functools import reduce

K = TypeVar('K', bound=Hashable)
V1 = TypeVar('V1')
V2 = TypeVar('V2')
V3 = TypeVar('V3')
V4 = TypeVar('V4')
V5 = TypeVar('V5')

def keyset(m: Mapping[K, Any]) -> set[K]:
    return set(m.keys())

# 按常用参数个数逐一定义精确重载
@overload
def inner_join(m1: Mapping[K, V1]) -> dict[K, tuple[V1]]: ...
@overload
def inner_join(m1: Mapping[K, V1], m2: Mapping[K, V2]) -> dict[K, tuple[V1, V2]]: ...
@overload
def inner_join(m1: Mapping[K, V1], m2: Mapping[K, V2], m3: Mapping[K, V3]) -> dict[K, tuple[V1, V2, V3]]: ...
@overload
def inner_join(m1: Mapping[K, V1], m2: Mapping[K, V2], m3: Mapping[K, V3], m4: Mapping[K, V4]) -> dict[K, tuple[V1, V2, V3, V4]]: ...
@overload
def inner_join(m1: Mapping[K, V1], m2: Mapping[K, V2], m3: Mapping[K, V3], m4: Mapping[K, V4], m5: Mapping[K, V5]) -> dict[K, tuple[V1, V2, V3, V4, V5]]: ...
# 兜底重载，处理超过5个参数的场景，类型精度退化为Any
@overload
def inner_join(*ms: Mapping[K, Any]) -> dict[K, tuple[Any, ...]]: ...

def inner_join(*ms: Mapping[K, Any]) -> dict[K, tuple[Any, ...]]:
    keys = reduce(set.intersection, map(keyset, ms)) if ms else set()
    return {k: tuple(m[k] for m in ms) for k in keys}

这个方案的缺点是只能覆盖提前写了重载的参数个数，超出个数的场景无法精确推导，但绝大多数业务场景下用到的join参数不会超过5个，足够日常使用。

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