如何为柯里化函数添加符合mypy要求的类型注解？

问题描述 #

使用Concatenate和ParamSpec实现柯里化函数时，mypy在调用柯里化后的函数（如f(1)）时报错，但相同代码在pyright中可正常识别，需要修改类型注解以通过mypy检查。

原问题代码 #

from typing import TypeVar, ParamSpec, Concatenate
from typing import Callable as Fn, reveal_type

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


def curry(f: Fn[Concatenate[T, P], R]) -> Fn[[T], Fn[P, R]]:
    """An attempt at currying."""

    def outer(x: T) -> Fn[P, R]:
        def inner(*args: P.args, **kwargs: P.kwargs) -> R:
            return f(x, *args, **kwargs)

        return inner

    return outer


@curry
def f(x: T, y: int) -> T:
    """Test function."""
    return x


def g(x: T, /) -> Fn[[int], T]:
    """Test function."""
    return lambda _: x

reveal_type(f)
reveal_type(g)
reveal_type(f(1)) # <- this fails in mypy
reveal_type(g(1))

问题原因 #

mypy对装饰器中泛型与ParamSpec的结合处理逻辑与pyright不同，会过早绑定泛型变量T，导致调用柯里化后的函数时无法正确推断后续参数和返回值类型。

解决方案 #

方案1：改用显式泛型函数定义（Python 3.12+） #

使用PEP 695引入的泛型函数语法，明确标记目标函数为泛型，帮助mypy正确追踪类型变量：

from typing import TypeVar, ParamSpec, Concatenate
from typing import Callable as Fn, reveal_type

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


def curry(f: Fn[Concatenate[T, P], R]) -> Fn[[T], Fn[P, R]]:
    """An attempt at currying."""

    def outer(x: T) -> Fn[P, R]:
        def inner(*args: P.args, **kwargs: P.kwargs) -> R:
            return f(x, *args, **kwargs)

        return inner

    return outer


@curry
def f[T](x: T, y: int) -> T:  # 显式声明泛型
    """Test function."""
    return x


def g(x: T, /) -> Fn[[int], T]:
    """Test function."""
    return lambda _: x

reveal_type(f)          # 推断为: Callable[[T], Callable[[int], T]]
reveal_type(g)          # 推断为: Callable[[T], Callable[[int], T]]
reveal_type(f(1))       # 推断为: Callable[[int], int]
reveal_type(g(1))       # 推断为: Callable[[int], int]

方案2：替换装饰器语法为显式调用 #

避免装饰器语法导致的泛型过早绑定，直接调用curry函数包装目标函数：

from typing import TypeVar, ParamSpec, Concatenate
from typing import Callable as Fn, reveal_type

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


def curry(f: Fn[Concatenate[T, P], R]) -> Fn[[T], Fn[P, R]]:
    """An attempt at currying."""

    def outer(x: T) -> Fn[P, R]:
        def inner(*args: P.args, **kwargs: P.kwargs) -> R:
            return f(x, *args, **kwargs)

        return inner

    return outer


# 不用装饰器，显式调用curry包装函数
def f(x: T, y: int) -> T:
    """Test function."""
    return x
f = curry(f)


def g(x: T, /) -> Fn[[int], T]:
    """Test function."""
    return lambda _: x

reveal_type(f)          # 推断为: Callable[[T], Callable[[int], T]]
reveal_type(g)          # 推断为: Callable[[T], Callable[[int], T]]
reveal_type(f(1))       # 推断为: Callable[[int], int]
reveal_type(g(1))       # 推断为: Callable[[int], int]

方案3：使用泛型类装饰器 #

通过定义泛型类作为装饰器，让mypy更清晰地追踪泛型变量的生命周期：

from typing import TypeVar, ParamSpec, Concatenate
from typing import Callable as Fn, reveal_type, Generic

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


class curry(Generic[T, P, R]):
    def __call__(self, f: Fn[Concatenate[T, P], R]) -> Fn[[T], Fn[P, R]]:
        def outer(x: T) -> Fn[P, R]:
            def inner(*args: P.args, **kwargs: P.kwargs) -> R:
                return f(x, *args, **kwargs)
            return inner
        return outer

curry = curry()


@curry
def f(x: T, y: int) -> T:
    """Test function."""
    return x


def g(x: T, /) -> Fn[[int], T]:
    """Test function."""
    return lambda _: x

reveal_type(f)          # 推断为: Callable[[T], Callable[[int], T]]
reveal_type(g)          # 推断为: Callable[[T], Callable[[int], T]]
reveal_type(f(1))       # 推断为: Callable[[int], int]
reveal_type(g(1))       # 推断为: Callable[[int], int]

以上三种方案均可解决mypy报错问题，可根据项目的Python版本和代码风格选择合适的方案。

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