diff --git a/packages/typestats-site/src/typestats_site/_uv.py b/packages/typestats-site/src/typestats_site/_uv.py index 203335ddc..3398a8a3e 100644 --- a/packages/typestats-site/src/typestats_site/_uv.py +++ b/packages/typestats-site/src/typestats_site/_uv.py @@ -1,5 +1,6 @@ import logging import os +import re import shutil import subprocess from typing import Final @@ -121,14 +122,59 @@ async def _is_top_level_module(p: anyio.Path) -> bool: return p.suffix in {".py", ".pyi"} and p.stem.isidentifier() -async def discover_packages(site_packages: StrPath, /) -> tuple[str, ...]: - """Return absolute paths of top-level packages/modules in *site_packages*. +def _normalize_dist(name: str) -> str: + """PEP 503 normalized distribution name.""" + return re.sub(r"[-_.]+", "-", name).lower() - Includes both package dirs (with `__init__.py[i]`) and single-file modules - (e.g. `six.py`). Falls back to *site_packages* itself when nothing matches. + +def _import_name(entry: str) -> str: + """Import name of a top-level `site-packages` entry (a `.py[i]` file's stem).""" + return re.sub(r"\.pyi?$", "", entry) + + +async def _dist_modules(sp: anyio.Path, dist_name: str) -> set[str] | None: + """Import names installed by *dist_name* (from its `RECORD`), or `None`.""" + target = _normalize_dist(dist_name) + async for child in sp.iterdir(): + if child.suffix != ".dist-info": + continue + + if _normalize_dist(child.stem.split("-", 1)[0]) != target: + continue + + record = child / "RECORD" + if not await record.exists(): + return None + + names: set[str] = set() + for line in (await record.read_text()).splitlines(): + top = line.split(",", 1)[0].split("/", 1)[0] + if not top or top.startswith(".") or top.endswith((".dist-info", ".data")): + continue + names.add(_import_name(top)) + return names + + return None + + +async def discover_packages( + site_packages: StrPath, + /, + dist_name: str | None = None, +) -> tuple[str, ...]: + """Absolute paths of top-level modules in *site_packages*. + + With *dist_name*, only that distribution's own modules (not its installed + dependencies) are returned. Falls back to *site_packages* when empty. """ sp = await anyio.Path(site_packages).resolve() - found = [str(p) async for p in sp.iterdir() if await _is_top_level_module(p)] + names = await _dist_modules(sp, dist_name) if dist_name else None + found = [ + str(p) + async for p in sp.iterdir() + if await _is_top_level_module(p) + and (names is None or _import_name(p.name) in names) + ] return tuple(found) or (str(sp),) diff --git a/packages/typestats-site/src/typestats_site/collect.py b/packages/typestats-site/src/typestats_site/collect.py index 4385b63e2..58b39d348 100644 --- a/packages/typestats-site/src/typestats_site/collect.py +++ b/packages/typestats-site/src/typestats_site/collect.py @@ -109,7 +109,7 @@ async def collect_version(self, version: Version, out: Path) -> bool: self.base_available = await available_versions(self.client, detected) pypi = PypiInfo.from_file_detail(self.eligible[version]) - pyrefly_paths = await discover_packages(sp) + pyrefly_paths = await discover_packages(sp, dist_name=project.name) if base_name := self.base_name: assert self.base_available is not None diff --git a/packages/typestats-site/src/typestats_site/from_project.py b/packages/typestats-site/src/typestats_site/from_project.py index dcb86f889..4f233dcb5 100644 --- a/packages/typestats-site/src/typestats_site/from_project.py +++ b/packages/typestats-site/src/typestats_site/from_project.py @@ -75,7 +75,7 @@ async def from_project( base_version=str(base_ver), exclude=project.exclude, pypi=PypiInfo.from_file_detail(dist_file), - pyrefly_paths=await discover_packages(sp), + pyrefly_paths=await discover_packages(sp, dist_name=project.name), ), ) @@ -86,6 +86,6 @@ async def from_project( FromPathOptions( exclude=project.exclude, pypi=PypiInfo.from_file_detail(dist_file), - pyrefly_paths=await discover_packages(sp), + pyrefly_paths=await discover_packages(sp, dist_name=project.name), ), ) diff --git a/packages/typestats-site/tests/test_uv.py b/packages/typestats-site/tests/test_uv.py index 339ea9405..80235ce98 100644 --- a/packages/typestats-site/tests/test_uv.py +++ b/packages/typestats-site/tests/test_uv.py @@ -312,6 +312,36 @@ async def test_returns_absolute_paths_from_relative_input( assert result == (str(pkg.resolve()),) assert Path(result[0]).is_absolute() + @pytest.mark.parametrize( + ("module", "distinfo", "dist_name"), + [ + ("mypkg", "mypkg-1.0.0.dist-info", "mypkg"), + ("my_mod.py", "My_Mod-2.0.dist-info", "My-Mod"), + ], + ids=["package", "normalized-file"], + ) + async def test_scopes_to_dist_modules( + self, tmp_path: Path, module: str, distinfo: str, dist_name: str + ) -> None: + """With dist_name, only that distribution's own modules are returned.""" + target = tmp_path / module + if module.endswith(".py"): + target.write_text("") + record = f"{module},,\n" + else: + target.mkdir() + (target / "__init__.py").write_text("") + record = f"{module}/__init__.py,,\n" + + dep = tmp_path / "dep" # installed dependency, not part of the distribution + dep.mkdir() + (dep / "__init__.py").write_text("") + (tmp_path / distinfo).mkdir() + (tmp_path / distinfo / "RECORD").write_text(record) + + result = await discover_packages(tmp_path, dist_name=dist_name) + assert result == (str(target.resolve()),) + async def test_mixed_package_and_module(self, tmp_path: Path) -> None: pkg = tmp_path / "mypkg" pkg.mkdir()