如何通过装饰器实现Python函数重复重载，保持代码简洁？

核心概要（TLDR） #

我希望通过指定函数重载来满足类型检查器的要求，但存在大量重复的函数签名，希望实现代码复用（DRY原则）。

示例与现有代码 #

我有多个函数，它们的第一个参数为int或float，且返回值类型与首参数一致。为让编辑器识别这种一一对应的类型关系，我手动添加了重载定义：

from typing import overload

@overload
def add_two(x: int) -> int: ...

@overload
def add_two(x: float) -> float: ...

def add_two(x: int | float) -> int | float:
    return x + 2

这种方式可行，但问题在于我有大量类似函数，例如：

def add_two(x: int | float) -> int | float:
    return x + 2

def square(x: int | float) -> int | float:
    return x ** 2

def multiply_by_int_add_int(x: int | float, y: int, z: int) -> int | float:
    return x * y + z

(...)

若为每个函数都添加上述重载定义，会产生大量重复代码。

问题 #

是否可以编写一个装饰器，将其应用到每个函数定义上以实现上述重载逻辑？

注意：

• 这些函数的共同点是：第一个参数的类型决定了返回值的类型。
• 其余参数的数量和类型不固定。

当前尝试 #

import functools
from typing import Callable, ParamSpec, TypeVar, overload

T = TypeVar("T")
P = ParamSpec("P")


def int2int_and_float2float(fn: Callable[P, T]) -> Callable[P, T]:

    @overload
    def wrapped(x: int, *args, **kwargs) -> int: ...

    @overload
    def wrapped(x: float, *args, **kwargs) -> float: ...

    @functools.wraps(fn)
    def wrapped(*args: P.args, **kwargs: P.kwargs):
        return fn(*args, **kwargs)

    return wrapped

@int2int_and_float2float
def add_two(x: int | float) -> int | float:
    """Add two"""
    return x * 2

该尝试未生效：add_two(3)的返回值未被识别为整数。

我忽略了哪些要点？

（需支持Python >=3.10版本）

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