如何为MagicMock自动生成类型提示以实现IDE自动补全

你当前手动编写TypingMockA这类类型提示类的方案虽然能实现自动补全，但存在重复劳动的问题——当目标类方法较多时，手动维护类型提示会非常繁琐。以下是几种更高效的解决方案：

方案1：Python 3.10+ 交集类型（推荐） #

利用Python 3.10引入的交集类型（&），让mock对象同时拥有目标类的方法提示和MagicMock的断言方法提示，无需手动编写类型类：

from typing import TypeVar, cast
from unittest.mock import MagicMock

T = TypeVar('T')

def auto_typing_mock(spec: type[T], **kwargs) -> T & MagicMock:
    return cast(T & MagicMock, MagicMock(spec=spec, **kwargs))

# 使用示例
class A:
    def f(self, a: int) -> int:
        print("f")
        return a

mock = auto_typing_mock(A)
mock.f(1)  # IDE自动补全f方法的参数、返回值
mock.f.assert_called_once()  # IDE自动补全MagicMock的断言方法

交集类型会告诉类型系统：这个mock对象同时是目标类A的实例和MagicMock的实例，因此IDE能同时识别两类成员。

方案2：泛型类实现（兼容Python 3.8+） #

如果使用Python 3.8/3.9（不支持交集类型），可以通过泛型类结合cast实现类似效果：

from typing import TypeVar, Generic, cast
from unittest.mock import MagicMock

T = TypeVar('T')

class AutoTypingMagicMock(Generic[T], MagicMock):
    @classmethod
    def create(cls, spec: type[T], **kwargs) -> 'AutoTypingMagicMock[T]':
        return cast(AutoTypingMagicMock[T], cls(spec=spec, **kwargs))

# 使用示例
class A:
    def f(self, a: int) -> int:
        print("f")
        return a

mock = AutoTypingMagicMock.create(A)
mock.f(1)  # IDE自动补全f方法
mock.f.assert_called_once()  # IDE自动补全断言方法

泛型类让IDE识别目标类的类型信息，同时继承MagicMock保留断言方法的提示。

方案3：自动生成类型提示类（静态代码生成） #

如果想完全复用手动编写类型类的思路，但不想手动编写，可以通过脚本自动生成对应目标类的类型提示代码：

import inspect
from typing import get_type_hints

def generate_typing_mock_class(target_cls: type) -> str:
    class_name = f"TypingMock{target_cls.__name__}"
    methods = []
    # 遍历目标类的公开方法
    for name, method in inspect.getmembers(target_cls, predicate=inspect.isfunction):
        if name.startswith('_'):
            continue
        hints = get_type_hints(method)
        # 整理方法参数和返回值类型
        args = []
        for param_name, param_type in hints.items():
            if param_name == 'return':
                continue
            args.append(f"{param_name}: {param_type.__name__}")
        return_type = hints.get('return', 'None').__name__
        methods.append(f"    def {name}({', '.join(args)}) -> {return_type}: ...")
    
    return f"class {class_name}:\n" + "\n".join(methods)

# 生成类A的类型提示类代码
class A:
    def f(self, a: int) -> int:
        print("f")
        return a

print(generate_typing_mock_class(A))

运行脚本会输出：

class TypingMockA:
    def f(self, a: int) -> int: ...

你可以将生成的代码复制到项目中使用，或集成到构建流程自动生成。

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