如何在Python3中实现类型安全的partial_apply函数

这个问题我之前也碰到过，想实现一个类型安全的偏应用函数，直接用subset_of、remainder_of这类自定义类型提示肯定行不通，不过Python的typing模块（配合typing_extensions）已经提供了足够的工具来解决这个问题，下面我给你一步步说清楚怎么实现：

核心思路 #

我们需要利用ParamSpec捕获原始函数的完整参数签名，再通过Concatenate把已传入的参数和剩余待传入的参数拼接起来，让类型检查器能准确推断出返回函数的参数类型。同时用Unpack来拆分参数，保证参数传递的类型安全。

完整实现代码 #

下面是支持位置参数偏应用的版本，兼容Python 3.10+（如果是更早版本，需要安装typing_extensions模块并从其中导入Concatenate和Unpack）：

from typing import Callable, ParamSpec, TypeVar, Unpack
from typing_extensions import Concatenate

# 定义参数规格和返回值类型变量
P = ParamSpec('P')
R = TypeVar('R')
# 用于捕获已传入的位置参数类型
PartialArgs = TypeVar('PartialArgs')

def partial_apply(
    fn: Callable[Concatenate[PartialArgs, P], R],
    *args: PartialArgs
) -> Callable[P, R]:
    """类型安全的偏应用函数：固定部分参数，返回接受剩余参数的新函数"""
    def wrapper(*wrapper_args: Unpack[P.args], **wrapper_kwargs: Unpack[P.kwargs]) -> R:
        # 拼接已传入的参数和新参数，调用原始函数
        return fn(*args, *wrapper_args, **wrapper_kwargs)
    return wrapper

# 测试用例
def add(a: int, b: int) -> int:
    return a + b

# 固定第一个参数为1，返回接受第二个int参数的函数
add_1 = partial_apply(add, 1)
# 类型检查器会推断add_1的类型是Callable[[int], int]
print(add_1(2))  # 输出3

# 如果传入错误类型的参数，类型检查器（比如mypy）会报错
# add_1("2")  # 提示：Argument 1 has incompatible type "str"; expected "int"

支持关键字参数的版本 #

如果需要同时支持位置参数和关键字参数的偏应用，可以调整代码如下：

from typing import Callable, ParamSpec, TypeVar, Unpack
from typing_extensions import Concatenate

P = ParamSpec('P')
R = TypeVar('R')
PartialArgs = TypeVar('PartialArgs')
PartialKwargs = TypeVar('PartialKwargs')

def partial_apply(
    fn: Callable[Concatenate[PartialArgs, P], R],
    *args: PartialArgs,
    **kwargs: PartialKwargs
) -> Callable[P, R]:
    def wrapper(*wrapper_args: Unpack[P.args], **wrapper_kwargs: Unpack[P.kwargs]) -> R:
        return fn(*args, *wrapper_args, **kwargs, **wrapper_kwargs)
    return wrapper

# 测试关键字参数偏应用
def greet(name: str, greeting: str = "Hello") -> str:
    return f"{greeting}, {name}!"

# 固定greeting参数为"Hi"，返回接受name参数的函数
greet_hi = partial_apply(greet, greeting="Hi")
print(greet_hi("Alice"))  # 输出"Hi, Alice!"

版本兼容性说明 #

• Python 3.10+：Concatenate和Unpack已经内置在typing模块中，可以直接导入使用。
• Python 3.9及以下：需要先安装typing_extensions（pip install typing_extensions），然后从typing_extensions中导入Concatenate和Unpack。

这种实现方式完全符合类型安全要求，类型检查器（比如mypy、pyright）会自动验证参数类型是否匹配，帮你提前发现错误。

备注：内容来源于stack exchange，提问作者luochen1990