如何实现输入输出类型一一对应的Python函数自动重载？

我有一类函数，它们的返回类型与单个输入的类型存在严格的一一对应关系。比如下面这个函数：

from typing import reveal_type

def func(x: str | list[str]) -> int | list[int]:
    """
    Some docstring that should show up.
    """
    if isinstance(x, str):
        return 0
    return list(range(len(x)))

y = func(["hi", "hello"])
reveal_type(y)  # 类型显示为 int | list[int]

直接使用联合类型标注输入输出时，IDE或类型检查器会把y的类型推断为int | list[int]，但实际上传入list[str]时返回的必然是list[int]，类型提示不够精确。

为了获得精确的类型提示，我们可以手动用@overload装饰器实现重载：

from typing import overload, reveal_type

@overload
def func(x: str) -> int: ...

@overload
def func(x: list[str]) -> list[int]: ...

def func(x: str | list[str]) -> int | list[int]:
    """
    Some docstring that should show up.
    """
    if isinstance(x, str):
        return 0
    return list(range(len(x)))

y = func(["hi", "hello"])
reveal_type(y)  # 类型显示为 list[int]

这样y的类型就能被精确推断为list[int]了，但我有很多这类函数，手动写重载会非常繁琐。我希望能有一种更简洁的方式，比如假设存在UnionMatch类型：

from typing import UnionMatch  # 仅为示例，实际不存在

def func(x: UnionMatch[str, list[str]]) -> UnionMatch[int, list[int]]:
    """
    Some docstring that should show up.
    """
    if isinstance(x, str):
        return 0
    return list(range(len(x)))

无需手动编写@overload，就能让类型检查器自动对应输入输出的类型关系。

现有方案的问题 #

有人针对两种类型的场景提供了一个基于泛型类的解决方案，但存在两个致命问题：

• 函数的文档字符串会丢失，IDE中悬停查看时只能看到泛型类的类型信息，看不到原函数的文档；
• 仅支持两种类型的对应，无法扩展到三种及以上的类型场景。

该方案的代码如下：

from typing import (
    reveal_type,
    Callable,
    Concatenate,
    Generic,
    overload,
    ParamSpec,
    TypeVar,
)

_C1 = TypeVar("_C1")
_C2 = TypeVar("_C2")
_P = ParamSpec("_P")
_T1 = TypeVar("_T1")
_T2 = TypeVar("_T2")

class FuncTemplate(Generic[_C1, _C2, _P, _T1, _T2]):

    def __init__(self, func: Callable[Concatenate[_C1 | _C2, _P], _T1 | _T2]) -> None:
        self.func = func
        # self.__call__.__doc__ = func.__doc__
        # 此写法无效，会报错：AttributeError: attribute '__doc__' of 'method' objects is not writable

    @overload
    def __call__(
        self,
        common_input: _C1,
        *args: _P.args,
        **kwargs: _P.kwargs,
    ) -> _T1: ...

    @overload
    def __call__(
        self,
        common_input: _C2,
        *args: _P.args,
        **kwargs: _P.kwargs,
    ) -> _T2: ...

    def __call__(
        self,
        common_input: _C1 | _C2,
        *args: _P.args,
        **kwargs: _P.kwargs,
    ) -> _T1 | _T2:
        return self.func(common_input, *args, **kwargs)

@FuncTemplate
def func(x: str | list[str]) -> int | list[int]:
    """
    Some docstring that should show up.
    """
    if isinstance(x, str):
        return 0
    return list(range(len(x)))

y = func(["hi", "hello"])
reveal_type(y)  # list[int]

z = func("hey")
reveal_type(z)  # int

在IDE中悬停查看func时，显示的内容是(function) func: FuncTemplate[str, list[str], (), int, list[int]]，完全看不到原函数的文档字符串。

请问如何实现支持多类型对应且保留文档字符串的自动重载功能？

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