Refactor/unify stub discovery (#16021)

* Add a StubFile class with sub-classes StdlibStubFile and ThirdPartyStubFile to describe stubs files.
* Add a stdlib_stubs function to collect the stdlib stub files for a specific Python version.
* Add a third_party_stubs function to collect all third-party stub files or the stub files of a specific distribution.
* Add path_stubs to collect all stub files in a certain path.
* Use the new functions in mypy_test.py, ty_test.py, and stubsabot.py.

Author: srittau

lib/ts_utils/stubs.py

@@ -0,0 +1,80 @@
+"""Stub file discovery."""
+
+from functools import cached_property
+from pathlib import Path
+
+from ts_utils.paths import STDLIB_PATH, STUBS_PATH, TESTS_DIR, distribution_path
+from ts_utils.utils import parse_stdlib_versions_file
+
+
+class StubFile:
+    """Base class for stub files."""
+
+    def __init__(self, path: Path) -> None:
+        self.path = path
+
+    def __fspath__(self) -> str:
+        return self.path.__fspath__()
+
+    def __str__(self) -> str:
+        return str(self.path)
+
+    @cached_property
+    def module_name(self) -> str:
+        return ".".join(self.module_parts)
+
+    @cached_property
+    def module_parts(self) -> tuple[str, ...]:
+        raise NotImplementedError
+
+
+class StdlibStubFile(StubFile):
+    """A stdlib stub file."""
+
+    @cached_property
+    def module_parts(self) -> tuple[str, ...]:
+        relative = self.path.relative_to(STDLIB_PATH)
+        parts = list(relative.parts[:-1])
+        if relative.name != "__init__.pyi":
+            parts.append(relative.stem)
+        return tuple(parts)
+
+
+class ThirdPartyStubFile(StubFile):
+    """A third-party stub file."""
+
+    @cached_property
+    def upstream_distribution(self) -> str:
+        return self.path.relative_to(STUBS_PATH).parts[0]
+
+    @cached_property
+    def module_parts(self) -> tuple[str, ...]:
+        relative = self.path.relative_to(STUBS_PATH)
+        parts = list(relative.parts[1:-1])
+        if relative.name != "__init__.pyi":
+            parts.append(relative.stem)
+        return tuple(parts)
+
+
+def stdlib_stubs(version: str) -> list[StdlibStubFile]:
+    """Return the stdlib stubs available for the requested Python version."""
+    module_versions = parse_stdlib_versions_file()
+    stubs = (StdlibStubFile(path) for path in path_stubs(STDLIB_PATH))
+    return [stub for stub in stubs if module_versions.is_supported(stub.module_name, version)]
+
+
+def third_party_stubs(distribution: str | None = None) -> list[ThirdPartyStubFile]:
+    """Return third-party stubs.
+
+    If distribution is None, return all third-party stubs. Otherwise,
+    return only stubs for the given distribution.
+    """
+    stub_path = distribution_path(distribution) if distribution else STUBS_PATH
+    return [ThirdPartyStubFile(path) for path in path_stubs(stub_path)]
+
+
+def path_stubs(path: Path) -> list[Path]:
+    """Return paths to all stub files in a certain path."""
+    if path.is_file():
+        return [path] if path.suffix == ".pyi" and TESTS_DIR not in path.parts else []
+    return sorted(p for p in path.rglob("*.pyi") if TESTS_DIR not in p.parts)

scripts/stubsabot.py

@@ -38,6 +38,7 @@
 
 from ts_utils.metadata import ObsoleteMetadata, StubMetadata, read_metadata, update_metadata
 from ts_utils.paths import PYRIGHT_CONFIG, STUBS_PATH, distribution_path
+from ts_utils.stubs import third_party_stubs
 
 TYPESHED_OWNER = "python"
 TYPESHED_API_URL = f"https://api.github.com/repos/{TYPESHED_OWNER}/typeshed"
@@ -534,7 +535,7 @@ async def analyze_github_diff(
     # https://docs.github.com/en/rest/commits/commits#compare-two-commits
     py_files: list[FileInfo] = [file for file in json_resp["files"] if Path(file["filename"]).suffix == ".py"]
     stub_path = distribution_path(distribution)
-    files_in_typeshed = set(stub_path.rglob("*.pyi"))
+    files_in_typeshed = {stub.path for stub in third_party_stubs(distribution)}
     py_files_stubbed_in_typeshed = [file for file in py_files if (stub_path / f"{file['filename']}i") in files_in_typeshed]
     return DiffAnalysis(py_files=py_files, py_files_stubbed_in_typeshed=py_files_stubbed_in_typeshed)
 
@@ -570,7 +571,7 @@ async def analyze_gitlab_diff(
         py_files.append(FileInfo(filename=filename, status=status, additions=additions, deletions=deletions))
 
     stub_path = distribution_path(distribution)
-    files_in_typeshed = set(stub_path.rglob("*.pyi"))
+    files_in_typeshed = {stub.path for stub in third_party_stubs(distribution)}
     py_files_stubbed_in_typeshed = [file for file in py_files if (stub_path / f"{file['filename']}i") in files_in_typeshed]
     return DiffAnalysis(py_files=py_files, py_files_stubbed_in_typeshed=py_files_stubbed_in_typeshed)
 

tests/mypy_test.py

@@ -22,14 +22,14 @@
 
 from ts_utils.metadata import PackageDependencies, get_recursive_requirements, read_metadata
 from ts_utils.mypy import MypyDistConf, mypy_configuration_from_distribution, temporary_mypy_config_file
-from ts_utils.paths import STDLIB_PATH, STUBS_PATH, TESTS_DIR, TS_BASE_PATH, distribution_path
+from ts_utils.paths import STDLIB_PATH, STUBS_PATH, TS_BASE_PATH, distribution_path
 from ts_utils.py315 import PY315_INCOMPATIBLE_RUNTIME_DEPENDENCIES
+from ts_utils.stubs import StubFile, stdlib_stubs, third_party_stubs
 from ts_utils.utils import (
     PYTHON_VERSION,
     colored,
     get_gitignore_spec,
     get_mypy_req,
-    parse_stdlib_versions_file,
     print_error,
     print_success_msg,
     spec_matches_path,
@@ -125,39 +125,28 @@ def log(args: TestConfig, *varargs: object) -> None:
         print(colored(" ".join(map(str, varargs)), "blue"))
 
 
-def match(path: Path, args: TestConfig) -> bool:
+def match(stub: StubFile, args: TestConfig) -> bool:
     for excluded_path in args.exclude:
-        if path == excluded_path:
-            log(args, path, "explicitly excluded")
+        if stub.path == excluded_path:
+            log(args, stub, "explicitly excluded")
             return False
-        if excluded_path in path.parents:
-            log(args, path, f'is in an explicitly excluded directory "{excluded_path}"')
+        if excluded_path in stub.path.parents:
+            log(args, stub, f'is in an explicitly excluded directory "{excluded_path}"')
             return False
     for included_path in args.filter:
-        if path == included_path:
-            log(args, path, "was explicitly included")
+        if stub.path == included_path:
+            log(args, stub, "was explicitly included")
             return True
-        if included_path in path.parents:
-            log(args, path, f'is in an explicitly included directory "{included_path}"')
+        if included_path in stub.path.parents:
+            log(args, stub, f'is in an explicitly included directory "{included_path}"')
             return True
     log_msg = (
         f'is implicitly excluded: was not in any of the directories or paths specified on the command line: "{args.filter!r}"'
     )
-    log(args, path, log_msg)
+    log(args, stub, log_msg)
     return False
 
 
-def add_files(files: list[Path], module: Path, args: TestConfig) -> None:
-    """Add all files in package or module represented by 'name' located in 'root'."""
-    if module.name.startswith("."):
-        return
-    if module.is_file() and module.suffix == ".pyi":
-        if match(module, args):
-            files.append(module)
-    else:
-        files.extend(sorted(file for file in module.rglob("*.pyi") if match(file, args)))
-
-
 class MypyResult(Enum):
     SUCCESS = 0
     FAILURE = 1
@@ -238,15 +227,14 @@ def run_mypy(
     return MypyResult.from_process_result(result)
 
 
-def add_third_party_files(distribution: str, files: list[Path], args: TestConfig, seen_dists: set[str]) -> None:
+def distribution_stub_files(distribution: str, args: TestConfig, seen_dists: set[str]) -> list[Path]:
     typeshed_reqs = get_recursive_requirements(distribution).typeshed_pkgs
     if distribution in seen_dists:
-        return
+        return []
     seen_dists.add(distribution)
     seen_dists.update(r.name for r in typeshed_reqs)
-    root = distribution_path(distribution)
-    for path in root.iterdir():
-        add_files(files, path, args)
+
+    return [stub.path for stub in third_party_stubs(distribution) if match(stub, args)]
 
 
 class TestResult(NamedTuple):
@@ -262,9 +250,8 @@ def test_third_party_distribution(
     Return a tuple, where the first element indicates mypy's return code
     and the second element is the number of checked files.
     """
-    files: list[Path] = []
     seen_dists: set[str] = set()
-    add_third_party_files(distribution, files, args, seen_dists)
+    files = distribution_stub_files(distribution, args, seen_dists)
     configurations = mypy_configuration_from_distribution(distribution)
 
     if not files and args.filter:
@@ -292,13 +279,7 @@ def test_third_party_distribution(
 
 
 def test_stdlib(args: TestConfig) -> TestResult:
-    files: list[Path] = []
-    for file in STDLIB_PATH.iterdir():
-        if file.name in ("VERSIONS", TESTS_DIR):
-            continue
-        add_files(files, file, args)
-
-    files = remove_modules_not_in_python_version(files, args.version)
+    files = [stub.path for stub in stdlib_stubs(args.version) if match(stub, args)]
 
     if not files:
         return TestResult(MypyResult.SUCCESS, 0)
@@ -309,27 +290,6 @@ 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]:
-    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)
-        if module_versions.is_supported(module_name, py_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":
-        parts.append(path.parts[-1].removesuffix(".pyi"))
-    return ".".join(parts)
-
-
 @dataclass
 class TestSummary:
     mypy_result: MypyResult = MypyResult.SUCCESS

tests/ty_test.py

@@ -9,7 +9,7 @@
 from pathlib import Path
 
 from ts_utils.paths import STDLIB_PATH, STUBS_PATH, TS_BASE_PATH
-from ts_utils.utils import parse_stdlib_versions_file
+from ts_utils.stubs import path_stubs, stdlib_stubs, third_party_stubs
 
 SUPPORTED_VERSIONS = ("3.10", "3.11", "3.12", "3.13", "3.14", "3.15")
 SUPPORTED_PLATFORMS = ("linux", "darwin", "win32")
@@ -19,34 +19,17 @@
 
 def stdlib_files(version: str) -> list[Path]:
     """Return the stdlib stubs available in the requested Python version."""
-    module_versions = parse_stdlib_versions_file()
-    files: list[Path] = []
+    # ty cannot resolve relative imports in the legacy distutils stubs.
+    return [stub.path for stub in stdlib_stubs(version) if stub.module_parts[0] != "distutils"]
 
-    for path in sorted(STDLIB_PATH.rglob("*.pyi")):
-        if "@tests" in path.parts:
-            continue
-        relative = path.relative_to(STDLIB_PATH)
-        # ty cannot resolve relative imports in the legacy distutils stubs.
-        if relative.parts[0] == "distutils":
-            continue
-        parts = list(relative.parts[:-1])
-        if relative.name != "__init__.pyi":
-            parts.append(relative.stem)
-        if module_versions.is_supported(".".join(parts), version):
-            files.append(path)
 
-    return files
-
-
-def _path_files(path: Path) -> list[Path]:
-    if path.is_file():
-        return [path] if path.suffix == ".pyi" and "@tests" not in path.parts else []
-    return sorted(p for p in path.rglob("*.pyi") if "@tests" not in p.parts)
+def third_party_files() -> list[Path]:
+    return [stub.path for stub in third_party_stubs() if stub.upstream_distribution not in EXCLUDED_STUBS]
 
 
 def _filter_files(files: list[Path], paths: list[Path], root: Path) -> list[Path]:
     selected_paths = [path if path.is_absolute() else TS_BASE_PATH / path for path in paths]
-    selected_files = {file for path in selected_paths if path.is_relative_to(root) for file in _path_files(path)}
+    selected_files = {file for path in selected_paths if path.is_relative_to(root) for file in path_stubs(path)}
     return [file for file in files if file in selected_files]
 
 
@@ -58,16 +41,12 @@ def main() -> int:
     parser.add_argument("--platform", choices=SUPPORTED_PLATFORMS, default="linux")
     args = parser.parse_args()
 
-    third_party_files = sorted(
-        path
-        for path in STUBS_PATH.rglob("*.pyi")
-        if "@tests" not in path.parts and path.relative_to(STUBS_PATH).parts[0] not in EXCLUDED_STUBS
-    )
     stdlib = stdlib_files(args.python_version)
+    third_party = third_party_files()
     if args.paths:
         stdlib = _filter_files(stdlib, args.paths, STDLIB_PATH)
-        third_party_files = _filter_files(third_party_files, args.paths, STUBS_PATH)
-    files = [*stdlib, *third_party_files]
+        third_party = _filter_files(third_party, args.paths, STUBS_PATH)
+    files = [*stdlib, *third_party]
     if not files:
         print("No stubs to check with ty.", flush=True)
         return 0