stubtest: use VERSIONS for submodules (#12083)

Closes #12062
Addresses feedback in #12066

Author: hauntsaninja

mypy/stubtest.py

@@ -1107,19 +1107,27 @@ def get_stub(module: str) -> Optional[nodes.MypyFile]:
     return _all_stubs.get(module)
 
 
-def get_typeshed_stdlib_modules(custom_typeshed_dir: Optional[str]) -> List[str]:
+def get_typeshed_stdlib_modules(
+    custom_typeshed_dir: Optional[str],
+    version_info: Optional[Tuple[int, int]] = None
+) -> List[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)
-    packages = set()
+    if version_info is None:
+        version_info = sys.version_info[0:2]
     # Typeshed's minimum supported Python 3 is Python 3.6
     if sys.version_info < (3, 6):
         version_info = (3, 6)
-    else:
-        version_info = sys.version_info[0:2]
-    for module, versions in stdlib_py_versions.items():
-        minver, maxver = versions
-        if version_info >= minver and (maxver is None or version_info <= maxver):
-            packages.add(module)
+
+    def exists_in_version(module: str) -> bool:
+        assert version_info is not None
+        parts = module.split(".")
+        for i in range(len(parts), 0, -1):
+            current_module = ".".join(parts[:i])
+            if current_module in stdlib_py_versions:
+                minver, maxver = stdlib_py_versions[current_module]
+                return version_info >= minver and (maxver is None or version_info <= maxver)
+        return False
 
     if custom_typeshed_dir:
         typeshed_dir = Path(custom_typeshed_dir)
@@ -1132,7 +1140,7 @@ def get_typeshed_stdlib_modules(custom_typeshed_dir: Optional[str]) -> List[str]
         if path.stem == "__init__":
             path = path.parent
         module = ".".join(path.relative_to(stdlib_dir).parts[:-1] + (path.stem,))
-        if module.split(".")[0] in packages:
+        if exists_in_version(module):
             modules.append(module)
     return sorted(modules)
 

mypy/test/teststubtest.py

@@ -983,12 +983,19 @@ def test_missing_stubs(self) -> None:
         assert "error: not_a_module failed to find stubs" in remove_color_code(output.getvalue())
 
     def test_get_typeshed_stdlib_modules(self) -> None:
-        stdlib = mypy.stubtest.get_typeshed_stdlib_modules(None)
+        stdlib = mypy.stubtest.get_typeshed_stdlib_modules(None, (3, 6))
         assert "builtins" in stdlib
         assert "os" in stdlib
         assert "os.path" in stdlib
         assert "asyncio" in stdlib
-        assert ("dataclasses" in stdlib) == (sys.version_info >= (3, 7))
+        assert "graphlib" not in stdlib
+        assert "formatter" in stdlib
+        assert "importlib.metadata" not in stdlib
+
+        stdlib = mypy.stubtest.get_typeshed_stdlib_modules(None, (3, 10))
+        assert "graphlib" in stdlib
+        assert "formatter" not in stdlib
+        assert "importlib.metadata" in stdlib
 
     def test_signature(self) -> None:
         def f(a: int, b: int, *, c: int, d: int = 0, **kwargs: Any) -> None: