diff --git a/mypy/stubtest.py b/mypy/stubtest.py index cb005c4bd52c..2f2b1943568c 100644 --- a/mypy/stubtest.py +++ b/mypy/stubtest.py @@ -1113,19 +1113,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) @@ -1138,7 +1146,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) diff --git a/mypy/test/teststubtest.py b/mypy/test/teststubtest.py index 6cc556004337..85f1aadb02c0 100644 --- a/mypy/test/teststubtest.py +++ b/mypy/test/teststubtest.py @@ -948,12 +948,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: