使用装饰器进行重复重载

Question

TLDR

我试图通过指定函数的重载来让类型检查器满意。我有很多特定的签名，我希望保持干燥。

示例和当前代码

我有许多函数将

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）

Answer 1

使函数通用，而不是使用重载。

T = TypeVar('T', int, float)

def add_two(x: T) -> T:
    return x + 2

def square(x: T) -> T:
    return x ** 2

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