如何通过mypy插件为支持懒加载的偏函数修饰器修正返回类型推断与参数检查逻辑

我完全懂你的痛点——你已经让mypy正确识别了两种场景的返回类型，但那个“缺少magic参数”的错误提示就像顽固的小尾巴，怎么甩都甩不掉。本质问题在于：mypy默认只会用原始函数的签名去校验所有调用，而你的修饰器实际上给函数赋予了两种合法的调用方式，所以我们需要让mypy知道这一点。

核心解决方案：给被修饰函数添加重载签名 #

我们可以通过mypy插件为被partial_decorator修饰的函数生成双重重载类型：

1. 第一种重载：保留原始函数的完整参数（包括magic），返回原函数的返回类型
2. 第二种重载：移除magic参数，返回PartialFunction类型

这样mypy会根据你调用时是否传入magic，自动匹配对应的重载，既不会报错参数缺失，也能正确推断返回类型。

修改后的插件代码 #

替换你原来的partial_plugin.py为以下内容：

from collections.abc import Callable
from typing import cast

from mypy.plugin import CallableTypeContext, Plugin, FunctionContext
from mypy.types import (
    CallableType,
    Instance,
    Overloaded,
    Type,
    get_proper_type,
)


def _adjust_callable_type(ctx: CallableTypeContext) -> Type:
    # 获取原始被修饰函数的可调用类型
    original_type = get_proper_type(ctx.type)
    if not isinstance(original_type, CallableType):
        return ctx.type

    # 找到magic参数的位置（处理位置参数和关键字参数）
    magic_param_idx = None
    for idx, name in enumerate(original_type.arg_names):
        if name == "magic":
            magic_param_idx = idx
            break
    if magic_param_idx is None:
        # 如果函数没有magic参数，直接返回原类型
        return ctx.type

    # 1. 生成第一个重载：原始完整签名，返回原类型
    overload_full = original_type

    # 2. 生成第二个重载：移除magic参数，返回PartialFunction类型
    new_arg_types = original_type.arg_types[:magic_param_idx] + original_type.arg_types[magic_param_idx+1:]
    new_arg_names = original_type.arg_names[:magic_param_idx] + original_type.arg_names[magic_param_idx+1:]
    new_kinds = original_type.arg_kinds[:magic_param_idx] + original_type.arg_kinds[magic_param_idx+1:]
    
    overload_partial = original_type.copy_modified(
        arg_types=new_arg_types,
        arg_names=new_arg_names,
        arg_kinds=new_kinds,
        ret_type=ctx.api.named_type("my_partial.PartialFunction"),
    )

    # 将两个重载组合成Overloaded类型
    return Overloaded([overload_full, overload_partial])


def _partial_function_hook_callback(ctx: FunctionContext) -> Type:
    # 保留原来的返回类型推断逻辑，作为兜底（可选，因为重载已经处理了返回类型）
    if "magic" in ctx.callee_arg_names:
        magic_index = ctx.callee_arg_names.index("magic")
        if not ctx.args[magic_index]:
            return ctx.api.named_type("my_partial.PartialFunction")
    return ctx.default_return_type


class PartialFunctionPlugin(Plugin):
    def get_callable_type_hook(self, fullname: str) -> Callable[[CallableTypeContext], Type] | None:
        # 针对partial_decorator修饰后的函数，调整其可调用类型为重载类型
        if fullname == "my_partial.partial_decorator":
            return _adjust_callable_type
        return None

    def get_function_hook(self, fullname: str) -> Callable[[FunctionContext], Type] | None:
        if fullname.startswith("my_partial.") and not fullname.endswith("partial_decorator"):
            return _partial_function_hook_callback
        return None


def plugin(version: str) -> type[PartialFunctionPlugin]:
    return PartialFunctionPlugin

代码解释 #

1. get_callable_type_hook：这个钩子会在partial_decorator应用到函数上时触发，我们在这里生成重载类型：

   • 先提取原始函数的参数列表，找到magic参数的位置
   • 生成第一个重载：完全保留原始函数的签名和返回类型
   • 生成第二个重载：移除magic参数，将返回类型替换为PartialFunction
   • 把两个重载组合成Overloaded类型，作为修饰后函数的最终类型

2. 保留原get_function_hook：作为兜底逻辑，确保返回类型的推断万无一失（其实重载已经能处理返回类型，但保留也没坏处）

测试效果 #

现在重新运行python -m mypy my_partial.py，你会看到：

my_partial.py:42: note: Revealed type is "builtins.str"
my_partial.py:46: note: Revealed type is "my_partial.PartialFunction"

那个烦人的Missing positional argument "magic"错误已经消失了！同时两个reveal_type的结果依然正确。

额外注意事项 #

• 如果你的magic参数是仅关键字参数（比如定义为*, magic: str | None），代码依然能正常工作，因为我们是通过参数名来定位的
• 如果被修饰的函数没有magic参数，插件会直接跳过处理，不会影响其他函数

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