如何用Python中不同数量的参数注释相互依赖的函数参数和返回类型(< 3.10)？

Question

编辑备注1：到目前为止，我已经找到了解决这个问题的佩普612，从Python3.10开始，介绍了typing.ParamSpec。因此，这个问题专门涉及Python3.9或更早版本。

编辑注释2：最初的示例太窄了，因为它的返回类型与参数类型完全匹配，但问题实际上是关于一个更一般的情况，即函数的参数签名是相同的，但返回类型是不同的。(也欢迎允许不同参数签名的更通用解决方案。)

问题：，我使用一个转换函数，它接收一个函数作为参数，并返回另一个函数作为其结果。传递的函数可以有任意数量和类型的参数(为了简单起见，我们可以坚持位置参数)，并且返回函数的调用方式与原始函数相同(参数的数目和类型相同)，返回类型取决于传递函数的返回类型(但不一定等于；例如，它可以是包含原始返回类型值的元组，也可以包含给定类型的另一个值)。

如何对转换函数进行注释，以反映传递函数和返回函数的签名的相关性？

一个简单的例子：

from typing import Callable


def transform(original_function: Callable) -> Callable:
    def new_function(*args):
        extra_payload = <do some calculation>
        return original_function(*args), extra_payload
    return new_function


def f(x: int, s: str) -> bool:
    ...


f(3, 'abc')  # this is a valid call
f('abc', 3)  # PyCharm warns about wrong argument types

# The goal is to have a warning on the following line:
transform(f)('abc', 3)  # no warning with the mere Callable annotations above

有没有办法让PyCharm知道transform(f)有与f相同的参数签名？

如果转换后的函数有固定数量的参数，我可以这样做，例如(假设有两个参数)：

from typing import TypeVar, Callable, Tuple


X = TypeVar('X')
Y = TypeVar('Y')
Z = TypeVar('Z')


def transform(original_function: Callable[[X, Y], Z]) -> Callable[[X, Y], Tuple[Z, <some type>]]:
   ...

，但是我的transform函数比它更通用，我将它用于具有不同数量的参数的函数，然后我不知道如何在Callable的第一个参数中指定它(上面有[X, Y] )。

(如何)在引入typing.ParamSpec之前可以这样做？

Answer 1

在Python3.10之前，似乎不可能完全优雅地做到这一点，而正是由于这个原因，Python3.10引入了ParamSpec (按照MisterMiyagi女士的评论)。

但是，可以通过对不同数量的参数使用重载(请参阅PEP 484的相关部分)来接近：

from typing import TypeVar, overload, Callable, Tuple


T = TypeVar('T')
T1 = TypeVar('T1')
T2 = TypeVar('T2')
T3 = TypeVar('T3')

Payload = <some type>


@overload
def transform(function: Callable[[], T]) -> Callable[[], Tuple[T, Payload]]:
    ...


@overload
def transform(function: Callable[[T1], T]) -> Callable[[T1], Tuple[T, Payload]]:
    ...


@overload
def transform(function: Callable[[T1, T2], T]) -> Callable[[T1, T2], Tuple[T, Payload]]:
    ...


@overload
def transform(function: Callable[[T1, T2, T3], T]) -> Callable[[T1, T2, T3], Tuple[T, Payload]]:
    ...


def transform(original_function):
    def new_function(*args):
        extra_payload = <do some calculation>
        return original_function(*args), extra_payload
    return new_function