如何在Python中扩展Pytest类型定义，支持自定义mark的类型检查？

解决Pytest自定义Mark无类型提示的问题 #

方法1：在插件代码里直接扩展类型（不用单独写stubs） #

在你自己的插件代码里，用类型检查专用逻辑扩展pytest的mark类型，完全不影响运行时：

# 你的插件文件，比如 my_pytest_plugin.py
import pytest
from typing import Callable, TYPE_CHECKING

if TYPE_CHECKING:
    # 继承原有的MarkDecorator，仅添加自定义标记的类型提示
    class ExtendedMark(pytest.MarkDecorator):
        def my_mark(self, arg1: str, arg2: str) -> Callable[[Callable], Callable]:
            ...

    # 给pytest.mark替换为扩展后的类型提示
    pytest.mark = ExtendedMark()

# 实现my_mark的实际逻辑
def my_mark(arg1: str, arg2: str) -> Callable[[Callable], Callable]:
    # 这里写你的标记逻辑，比如返回parametrize或自定义标记
    return pytest.mark.parametrize(...)

# 注册标记到pytest，避免运行时警告
def pytest_configure(config):
    config.addinivalue_line("markers", "my_mark(arg1, arg2): 这里写你的标记说明")

方法2：编写不覆盖原有类型的存根文件 #

别搞pytest-stubs这种会全局覆盖的库，直接在你的插件包里添加类型声明：

1. 在插件根目录创建空的py.typed文件，告诉mypy这个包支持类型检查。
2. 创建_types/pytest.pyi存根文件：

# _types/pytest.pyi
import pytest
from typing import Callable

# 继承原类型进行扩展，保留所有原生mark的类型提示
class ExtendedMark(pytest.MarkDecorator):
    def my_mark(self, arg1: str, arg2: str) -> Callable[[Callable], Callable]:
        ...

# 导出所有pytest原有内容，替换mark为扩展后的类型
__all__ = [*pytest.__all__, "mark"]
mark: ExtendedMark = pytest.mark

3. 在pyproject.toml中配置mypy，指定存根目录：

[tool.mypy]
strict = true
mypypath = ["_types"]

重点提醒 #

• 绝对不能直接覆盖pytest.mark的类，必须继承原有的MarkDecorator，否则原生的@pytest.mark.parametrize等标记会丢失类型提示。
• 用TYPE_CHECKING隔离类型代码，运行时不会执行这部分逻辑，不会破坏原有功能。
• 一定要在pytest_configure里注册标记，否则运行测试时pytest会抛出未知标记的警告。

这么处理后，编辑器会给@pytest.mark.my_mark提供参数提示，mypy也会检查参数类型错误，同时原生pytest的类型定义完全不受影响。

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