Add mypy plugin support to stubtest configuration (#13948)

huynguyengl99 · web-flow · commit a8fa1ab0d9db · 2025-05-11T14:45:30.000+02:00
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -229,6 +229,12 @@ This has the following keys:
   If not specified, stubtest is run only on `linux`.
   Only add extra OSes to the test
   if there are platform-specific branches in a stubs package.
+* `mypy_plugins` (default: `[]`): A list of Python modules to use as mypy plugins
+when running stubtest. For example: `mypy_plugins = ["mypy_django_plugin.main"]`
+* `mypy_plugins_config` (default: `{}`): A dictionary mapping plugin names to their
+configuration dictionaries for use by mypy plugins. For example:
+`mypy_plugins_config = {"django-stubs" = {"django_settings_module" = "@tests.django_settings"}}`
+
 
 `*_dependencies` are usually packages needed to `pip install` the implementation
 distribution.
diff --git a/lib/ts_utils/metadata.py b/lib/ts_utils/metadata.py
@@ -11,7 +11,7 @@
 from collections.abc import Mapping
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Annotated, Final, NamedTuple, final
+from typing import Annotated, Any, Final, NamedTuple, final
 from typing_extensions import TypeGuard
 
 import tomli
@@ -42,6 +42,10 @@ def _is_list_of_strings(obj: object) -> TypeGuard[list[str]]:
     return isinstance(obj, list) and all(isinstance(item, str) for item in obj)
 
 
+def _is_nested_dict(obj: object) -> TypeGuard[dict[str, dict[str, Any]]]:
+    return isinstance(obj, dict) and all(isinstance(k, str) and isinstance(v, dict) for k, v in obj.items())
+
+
 @functools.cache
 def _get_oldest_supported_python() -> str:
     with PYPROJECT_PATH.open("rb") as config:
@@ -71,6 +75,8 @@ class StubtestSettings:
     ignore_missing_stub: bool
     platforms: list[str]
     stubtest_requirements: list[str]
+    mypy_plugins: list[str]
+    mypy_plugins_config: dict[str, dict[str, Any]]
 
     def system_requirements_for_platform(self, platform: str) -> list[str]:
         assert platform in _STUBTEST_PLATFORM_MAPPING, f"Unrecognised platform {platform!r}"
@@ -93,6 +99,8 @@ def read_stubtest_settings(distribution: str) -> StubtestSettings:
     ignore_missing_stub: object = data.get("ignore_missing_stub", False)
     specified_platforms: object = data.get("platforms", ["linux"])
     stubtest_requirements: object = data.get("stubtest_requirements", [])
+    mypy_plugins: object = data.get("mypy_plugins", [])
+    mypy_plugins_config: object = data.get("mypy_plugins_config", {})
 
     assert type(skip) is bool
     assert type(ignore_missing_stub) is bool
@@ -104,6 +112,8 @@ def read_stubtest_settings(distribution: str) -> StubtestSettings:
     assert _is_list_of_strings(choco_dependencies)
     assert _is_list_of_strings(extras)
     assert _is_list_of_strings(stubtest_requirements)
+    assert _is_list_of_strings(mypy_plugins)
+    assert _is_nested_dict(mypy_plugins_config)
 
     unrecognised_platforms = set(specified_platforms) - _STUBTEST_PLATFORM_MAPPING.keys()
     assert not unrecognised_platforms, f"Unrecognised platforms specified for {distribution!r}: {unrecognised_platforms}"
@@ -124,6 +134,8 @@ def read_stubtest_settings(distribution: str) -> StubtestSettings:
         ignore_missing_stub=ignore_missing_stub,
         platforms=specified_platforms,
         stubtest_requirements=stubtest_requirements,
+        mypy_plugins=mypy_plugins,
+        mypy_plugins_config=mypy_plugins_config,
     )
 
 
@@ -179,6 +191,8 @@ def is_obsolete(self) -> bool:
         "ignore_missing_stub",
         "platforms",
         "stubtest_requirements",
+        "mypy_plugins",
+        "mypy_plugins_config",
     }
 }
 _DIST_NAME_RE: Final = re.compile(r"^[a-z0-9]([a-z0-9._-]*[a-z0-9])?$", re.IGNORECASE)
diff --git a/lib/ts_utils/mypy.py b/lib/ts_utils/mypy.py
@@ -6,7 +6,7 @@
 
 import tomli
 
-from ts_utils.metadata import metadata_path
+from ts_utils.metadata import StubtestSettings, metadata_path
 from ts_utils.utils import NamedTemporaryFile, TemporaryFileWrapper
 
 
@@ -50,14 +50,27 @@ def validate_configuration(section_name: str, mypy_section: dict[str, Any]) -> M
 
 
 @contextmanager
-def temporary_mypy_config_file(configurations: Iterable[MypyDistConf]) -> Generator[TemporaryFileWrapper[str]]:
+def temporary_mypy_config_file(
+    configurations: Iterable[MypyDistConf], stubtest_settings: StubtestSettings | None = None
+) -> Generator[TemporaryFileWrapper[str]]:
     temp = NamedTemporaryFile("w+")
     try:
         for dist_conf in configurations:
             temp.write(f"[mypy-{dist_conf.module_name}]\n")
             for k, v in dist_conf.values.items():
                 temp.write(f"{k} = {v}\n")
         temp.write("[mypy]\n")
+
+        if stubtest_settings:
+            if stubtest_settings.mypy_plugins:
+                temp.write(f"plugins = {'.'.join(stubtest_settings.mypy_plugins)}\n")
+
+            if stubtest_settings.mypy_plugins_config:
+                for plugin_name, plugin_dict in stubtest_settings.mypy_plugins_config.items():
+                    temp.write(f"[mypy.plugins.{plugin_name}]\n")
+                    for k, v in plugin_dict.items():
+                        temp.write(f"{k} = {v}\n")
+
         temp.flush()
         yield temp
     finally:
diff --git a/tests/README.md b/tests/README.md
@@ -196,6 +196,23 @@ that stubtest reports to be missing should necessarily be added to the stub.
 For some implementation details, it is often better to add allowlist entries
 for missing objects rather than trying to match the runtime in every detail.
 
+### Support for mypy plugins in stubtest
+
+For stubs that require mypy plugins to check correctly (such as Django), stubtest
+supports configuring mypy plugins through the METADATA.toml file. This allows stubtest to
+leverage type information provided by these plugins when validating stubs.
+
+To use this feature, add the following configuration to the `tool.stubtest` section in your METADATA.toml:
+
+```toml
+mypy_plugins = ["mypy_django_plugin.main"]
+mypy_plugins_config = { "django-stubs" = { "django_settings_module" = "@tests.django_settings" } }
+```
+
+For Django stubs specifically, you'll need to create a `django_settings.py` file in your `@tests` directory
+that contains the Django settings required by the plugin. This file will be referenced by the plugin
+configuration to properly validate Django-specific types during stubtest execution.
+
 ## typecheck\_typeshed.py
 
 Run using
diff --git a/tests/check_typeshed_structure.py b/tests/check_typeshed_structure.py
@@ -27,6 +27,10 @@
 # consistent CI runs.
 linters = {"mypy", "pyright", "pytype", "ruff"}
 
+ALLOWED_PY_FILES_IN_TESTS_DIR = {
+    "django_settings.py"  # This file contains Django settings used by the mypy_django_plugin during stubtest execution.
+}
+
 
 def assert_consistent_filetypes(
     directory: Path, *, kind: str, allowed: set[str], allow_nonidentifier_filenames: bool = False
@@ -81,7 +85,9 @@ def check_stubs() -> None:
 
 
 def check_tests_dir(tests_dir: Path) -> None:
-    py_files_present = any(file.suffix == ".py" for file in tests_dir.iterdir())
+    py_files_present = any(
+        file.suffix == ".py" and file.name not in ALLOWED_PY_FILES_IN_TESTS_DIR for file in tests_dir.iterdir()
+    )
     error_message = f"Test-case files must be in an `{TESTS_DIR}/{TEST_CASES_DIR}` directory, not in the `{TESTS_DIR}` directory"
     assert not py_files_present, error_message
 
diff --git a/tests/stubtest_third_party.py b/tests/stubtest_third_party.py
@@ -97,7 +97,7 @@ def run_stubtest(
             return False
 
         mypy_configuration = mypy_configuration_from_distribution(dist_name)
-        with temporary_mypy_config_file(mypy_configuration) as temp:
+        with temporary_mypy_config_file(mypy_configuration, stubtest_settings) as temp:
             ignore_missing_stub = ["--ignore-missing-stub"] if stubtest_settings.ignore_missing_stub else []
             packages_to_check = [d.name for d in dist.iterdir() if d.is_dir() and d.name.isidentifier()]
             modules_to_check = [d.stem for d in dist.iterdir() if d.is_file() and d.suffix == ".pyi"]
