Python使用Protocol与ParamSpec标注装饰器类型的mypy报错求解

• 实现了一个装饰器，接收函数作为入参，返回与原函数签名完全一致的函数。被装饰函数的第一个参数必须带有foo属性，装饰器会基于该属性值执行对应副作用操作。
• 在Python 3.10环境下结合Protocol和ParamSpec编写类型提示时，mypy检查抛出类型不兼容错误。

复现代码 #

import functools
from typing import Callable, Concatenate, ParamSpec, Protocol, TypeVar


class HasFoo(Protocol):
    foo: int


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


def print_foo(
    f: Callable[Concatenate[HasFoo, P], T]
) -> Callable[Concatenate[HasFoo, P], T]:
    """This is my decorator.

    It prints `x.foo` where `x` is the first argument of the
    decorated function.
    """

    @functools.wraps(f)
    def wrapped(has_foo: HasFoo, *args: P.args, **kwargs: P.kwargs) -> T:
        print(has_foo.foo)
        return f(has_foo, *args, **kwargs)

    return wrapped


class ReallyHasFoo:
    """Some actual implementation of `HasFoo`."""

    def __init__(self) -> None:
        self.foo = 0


@print_foo
def f(has_foo: ReallyHasFoo) -> None:
    pass

报错信息 #

执行mypy检查时抛出如下错误：

test.py:37: error: Argument 1 to "print_foo" has incompatible type "Callable[[ReallyHasFoo], None]"; expected "Callable[[HasFoo], None]"
test.py:37: note: This may be because "f" has arguments named: "has_foo"

Callable类型的参数是逆变的：如果标注要求传入的函数第一个参数为HasFoo类型，意味着这个函数必须能接受所有实现了HasFoo协议的实例。但示例中被装饰的函数f只接受ReallyHasFoo类型的参数，不满足入参要求，因此mypy会报错。

给装饰器的第一个参数单独定义一个绑定到HasFoo协议的类型变量，不再将参数类型固定写死为HasFoo，让类型检查器可以自动推断被装饰函数第一个参数的具体子类型。
修正后的完整代码如下：

import functools
from typing import Callable, Concatenate, ParamSpec, Protocol, TypeVar


class HasFoo(Protocol):
    foo: int


P = ParamSpec("P")
T = TypeVar("T")
# 定义绑定到HasFoo协议的类型变量，可匹配任意实现HasFoo的具体类型
F = TypeVar("F", bound=HasFoo)


def print_foo(
    f: Callable[Concatenate[F, P], T]
) -> Callable[Concatenate[F, P], T]:
    """This is my decorator.

    It prints `x.foo` where `x` is the first argument of the
    decorated function.
    """

    @functools.wraps(f)
    def wrapped(has_foo: F, *args: P.args, **kwargs: P.kwargs) -> T:
        print(has_foo.foo)
        return f(has_foo, *args, **kwargs)

    return wrapped


class ReallyHasFoo:
    """Some actual implementation of `HasFoo`."""

    def __init__(self) -> None:
        self.foo = 0


@print_foo
def f(has_foo: ReallyHasFoo) -> None:
    pass

修改后mypy检查可以正常通过，同时装饰器返回的函数会完整保留原函数的签名类型，不会丢失参数类型信息。

内容的提问来源于stack exchange，提问作者Roméo Després