如何为类实现的装饰器编写类型提示？PyCharm识别异常解惑

问题背景 #

先看一个经典的类实现装饰器示例：

class Decorator:
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        self.func(*args, **kwargs)

为了让该装饰器的__call__方法拥有与被装饰函数func完全一致的签名和类型提示，尝试了基于PEP 612规范的类型注解实现：

from typing import Callable, TypeVar, ParamSpec, Generic

PT = ParamSpec('PT')
RT = TypeVar('RT')


class Decorator(Generic[PT, RT]):
    def __init__(self, func: Callable[PT, RT]) -> None:
        self.func = func

    def __call__(self, *args: PT.args, **kwargs: PT.kwargs) -> RT:
        return self.func(*args, **kwargs)


@Decorator
def add(x: int, y: int) -> int:
    return x + y

但在PyCharm中调用add时，无法获取到正确的参数列表（调用提示显示参数为*args, **kwargs，而非预期的x: int, y: int），这是PyCharm的问题吗？

问题原因与解决方案 #

这确实是PyCharm对基于ParamSpec和Generic的类装饰器类型推断支持不完善导致的问题——代码本身符合PEP规范，在mypy等标准类型检查器中能正常识别签名和类型提示。

可以通过以下几种方式解决：

方案1：补充函数元数据 #

在装饰器类的__init__方法中使用functools.wraps复制原函数的元数据，帮助PyCharm识别正确签名：

from typing import Callable, TypeVar, ParamSpec, Generic
import functools

PT = ParamSpec('PT')
RT = TypeVar('RT')


class Decorator(Generic[PT, RT]):
    def __init__(self, func: Callable[PT, RT]) -> None:
        self.func = func
        functools.wraps(func)(self)  # 复制原函数的名称、文档、参数签名等元数据

    def __call__(self, *args: PT.args, **kwargs: PT.kwargs) -> RT:
        return self.func(*args, **kwargs)


@Decorator
def add(x: int, y: int) -> int:
    return x + y

方案2：改用函数式装饰器 #

如果不需要类装饰器的额外特性，函数式装饰器的类型推断在PyCharm中支持更完善：

from typing import Callable, TypeVar, ParamSpec, Generic

PT = ParamSpec('PT')
RT = TypeVar('RT')


def decorator(func: Callable[PT, RT]) -> Callable[PT, RT]:
    def wrapper(*args: PT.args, **kwargs: PT.kwargs) -> RT:
        return func(*args, **kwargs)
    return wrapper


@decorator
def add(x: int, y: int) -> int:
    return x + y

方案3：升级PyCharm版本 #

JetBrains在后续版本中逐步完善了对PEP 612的支持，升级到2024.1及以上版本的PyCharm，大概率可以直接解决这个类型推断问题。

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