mypy_test: Exclude sub-modules not in current Py version (#12352)

Author: srittau

tests/_utils.py

@@ -7,7 +7,7 @@
 from collections.abc import Iterable, Mapping
 from functools import lru_cache
 from pathlib import Path
-from typing import Any, Final, NamedTuple, Tuple
+from typing import Any, Dict, Final, NamedTuple, Tuple
 from typing_extensions import TypeAlias
 
 import pathspec
@@ -108,14 +108,14 @@ def get_mypy_req() -> str:
 # ====================================================================
 
 VersionTuple: TypeAlias = Tuple[int, int]
-
+SupportedVersionsDict: TypeAlias = Dict[str, Tuple[VersionTuple, VersionTuple]]
 
 VERSIONS_PATH = STDLIB_PATH / "VERSIONS"
 VERSION_LINE_RE = re.compile(r"^([a-zA-Z_][a-zA-Z0-9_.]*): ([23]\.\d{1,2})-([23]\.\d{1,2})?$")
 VERSION_RE = re.compile(r"^([23])\.(\d+)$")
 
 
-def parse_stdlib_versions_file() -> dict[str, tuple[VersionTuple, VersionTuple]]:
+def parse_stdlib_versions_file() -> SupportedVersionsDict:
     result: dict[str, tuple[VersionTuple, VersionTuple]] = {}
     with VERSIONS_PATH.open(encoding="UTF-8") as f:
         for line in f:

tests/mypy_test.py

@@ -24,7 +24,10 @@
 from _metadata import PackageDependencies, get_recursive_requirements, read_metadata
 from _utils import (
     PYTHON_VERSION,
+    STDLIB_PATH,
     TESTS_DIR,
+    SupportedVersionsDict,
+    VersionTuple,
     colored,
     get_gitignore_spec,
     get_mypy_req,
@@ -332,15 +335,12 @@ def test_third_party_distribution(
 
 def test_stdlib(args: TestConfig) -> TestResult:
     files: list[Path] = []
-    stdlib = Path("stdlib")
-    supported_versions = parse_stdlib_versions_file()
-    for name in os.listdir(stdlib):
-        if name in ("VERSIONS", TESTS_DIR) or name.startswith("."):
+    for file in STDLIB_PATH.iterdir():
+        if file.name in ("VERSIONS", TESTS_DIR) or file.name.startswith("."):
             continue
-        module = Path(name).stem
-        module_min_version, module_max_version = supported_versions[module]
-        if module_min_version <= tuple(map(int, args.version.split("."))) <= module_max_version:
-            add_files(files, (stdlib / name), args)
+        add_files(files, file, args)
+
+    files = remove_modules_not_in_python_version(files, args.version)
 
     if not files:
         return TestResult(MypyResult.SUCCESS, 0)
@@ -351,6 +351,38 @@ def test_stdlib(args: TestConfig) -> TestResult:
     return TestResult(result, len(files))
 
 
+def remove_modules_not_in_python_version(paths: list[Path], py_version: VersionString) -> list[Path]:
+    py_version_tuple = tuple(map(int, py_version.split(".")))
+    module_versions = parse_stdlib_versions_file()
+    new_paths: list[Path] = []
+    for path in paths:
+        if path.parts[0] != "stdlib" or path.suffix != ".pyi":
+            continue
+        module_name = stdlib_module_name_from_path(path)
+        min_version, max_version = supported_versions_for_module(module_versions, module_name)
+        if min_version <= py_version_tuple <= max_version:
+            new_paths.append(path)
+    return new_paths
+
+
+def stdlib_module_name_from_path(path: Path) -> str:
+    assert path.parts[0] == "stdlib"
+    assert path.suffix == ".pyi"
+    parts = list(path.parts[1:-1])
+    if path.parts[-1] != "__init__.pyi":
+        # TODO: Python 3.9+: Use removesuffix.
+        parts.append(path.parts[-1][:-4])
+    return ".".join(parts)
+
+
+def supported_versions_for_module(module_versions: SupportedVersionsDict, module_name: str) -> tuple[VersionTuple, VersionTuple]:
+    while "." in module_name:
+        if module_name in module_versions:
+            return module_versions[module_name]
+        module_name = ".".join(module_name.split(".")[:-1])
+    return module_versions[module_name]
+
+
 @dataclass
 class TestSummary:
     mypy_result: MypyResult = MypyResult.SUCCESS