From 0b303b53479897e24d57affef6a8cdfffbd08e3d Mon Sep 17 00:00:00 2001
From: Alex Waygood <Alex.Waygood@Gmail.com>
Date: Thu, 24 Aug 2023 00:47:22 +0100
Subject: [PATCH] stubtest: error if typeshed is missing modules from the
 stdlib (#15729)

We currently flag modules missing from third-party stubs in stubtest,
but don't do similarly for typeshed's stdlib stubs. This PR adds that
functionality for typeshed's stdlib stubs as well.
---
 mypy/stubtest.py | 83 ++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 74 insertions(+), 9 deletions(-)

diff --git a/mypy/stubtest.py b/mypy/stubtest.py
index b2506e6dcc02..d8a613034b3a 100644
--- a/mypy/stubtest.py
+++ b/mypy/stubtest.py
@@ -11,6 +11,7 @@
 import copy
 import enum
 import importlib
+import importlib.machinery
 import inspect
 import os
 import pkgutil
@@ -25,7 +26,7 @@
 from contextlib import redirect_stderr, redirect_stdout
 from functools import singledispatch
 from pathlib import Path
-from typing import Any, Generic, Iterator, TypeVar, Union
+from typing import AbstractSet, Any, Generic, Iterator, TypeVar, Union
 from typing_extensions import get_origin, is_typeddict
 
 import mypy.build
@@ -1639,7 +1640,7 @@ def get_stub(module: str) -> nodes.MypyFile | None:
 
 def get_typeshed_stdlib_modules(
     custom_typeshed_dir: str | None, version_info: tuple[int, int] | None = None
-) -> list[str]:
+) -> set[str]:
     """Returns a list of stdlib modules in typeshed (for current Python version)."""
     stdlib_py_versions = mypy.modulefinder.load_stdlib_py_versions(custom_typeshed_dir)
     if version_info is None:
@@ -1661,14 +1662,75 @@ def exists_in_version(module: str) -> bool:
         typeshed_dir = Path(mypy.build.default_data_dir()) / "typeshed"
     stdlib_dir = typeshed_dir / "stdlib"
 
-    modules = []
+    modules: set[str] = set()
     for path in stdlib_dir.rglob("*.pyi"):
         if path.stem == "__init__":
             path = path.parent
         module = ".".join(path.relative_to(stdlib_dir).parts[:-1] + (path.stem,))
         if exists_in_version(module):
-            modules.append(module)
-    return sorted(modules)
+            modules.add(module)
+    return modules
+
+
+def get_importable_stdlib_modules() -> set[str]:
+    """Return all importable stdlib modules at runtime."""
+    all_stdlib_modules: AbstractSet[str]
+    if sys.version_info >= (3, 10):
+        all_stdlib_modules = sys.stdlib_module_names
+    else:
+        all_stdlib_modules = set(sys.builtin_module_names)
+        python_exe_dir = Path(sys.executable).parent
+        for m in pkgutil.iter_modules():
+            finder = m.module_finder
+            if isinstance(finder, importlib.machinery.FileFinder):
+                finder_path = Path(finder.path)
+                if (
+                    python_exe_dir in finder_path.parents
+                    and "site-packages" not in finder_path.parts
+                ):
+                    all_stdlib_modules.add(m.name)
+
+    importable_stdlib_modules: set[str] = set()
+    for module_name in all_stdlib_modules:
+        if module_name in ANNOYING_STDLIB_MODULES:
+            continue
+
+        try:
+            runtime = silent_import_module(module_name)
+        except ImportError:
+            continue
+        else:
+            importable_stdlib_modules.add(module_name)
+
+        try:
+            # some stdlib modules (e.g. `nt`) don't have __path__ set...
+            runtime_path = runtime.__path__
+            runtime_name = runtime.__name__
+        except AttributeError:
+            continue
+
+        for submodule in pkgutil.walk_packages(runtime_path, runtime_name + "."):
+            submodule_name = submodule.name
+
+            # There are many annoying *.__main__ stdlib modules,
+            # and including stubs for them isn't really that useful anyway:
+            # tkinter.__main__ opens a tkinter windows; unittest.__main__ raises SystemExit; etc.
+            #
+            # The idlelib.* submodules are similarly annoying in opening random tkinter windows,
+            # and we're unlikely to ever add stubs for idlelib in typeshed
+            # (see discussion in https://github.com/python/typeshed/pull/9193)
+            if submodule_name.endswith(".__main__") or submodule_name.startswith("idlelib."):
+                continue
+
+            try:
+                silent_import_module(submodule_name)
+            # importing multiprocessing.popen_forkserver on Windows raises AttributeError...
+            except Exception:
+                continue
+            else:
+                importable_stdlib_modules.add(submodule_name)
+
+    return importable_stdlib_modules
 
 
 def get_allowlist_entries(allowlist_file: str) -> Iterator[str]:
@@ -1699,6 +1761,10 @@ class _Arguments:
     version: str
 
 
+# typeshed added a stub for __main__, but that causes stubtest to check itself
+ANNOYING_STDLIB_MODULES: typing_extensions.Final = frozenset({"antigravity", "this", "__main__"})
+
+
 def test_stubs(args: _Arguments, use_builtins_fixtures: bool = False) -> int:
     """This is stubtest! It's time to test the stubs!"""
     # Load the allowlist. This is a series of strings corresponding to Error.object_desc
@@ -1721,10 +1787,9 @@ def test_stubs(args: _Arguments, use_builtins_fixtures: bool = False) -> int:
                 "cannot pass both --check-typeshed and a list of modules",
             )
             return 1
-        modules = get_typeshed_stdlib_modules(args.custom_typeshed_dir)
-        # typeshed added a stub for __main__, but that causes stubtest to check itself
-        annoying_modules = {"antigravity", "this", "__main__"}
-        modules = [m for m in modules if m not in annoying_modules]
+        typeshed_modules = get_typeshed_stdlib_modules(args.custom_typeshed_dir)
+        runtime_modules = get_importable_stdlib_modules()
+        modules = sorted((typeshed_modules | runtime_modules) - ANNOYING_STDLIB_MODULES)
 
     if not modules:
         print(_style("error:", color="red", bold=True), "no modules to check")