TheAlgorithms/Python · #15374
pr_file_map.py: fetch PR files concurrently with asyncio
scripts/pr_file_map.py80 + / 9 −
@@ -31,7 +31,7 @@ Run status is also written to stderr with: - Number of PRs from `get_open_prs()`- - Number of file touches from `get_pr_files()` and distinct files touched+ - Number of file touches from `get_pr_files_async()` and distinct files touched - Number of existing and missing files Requirements: gh (GitHub CLI), authenticated (`gh auth login`)@@ -41,7 +41,9 @@ scripts/pr_file_map.py > report.md """ +import asyncio import json+import os import shutil import subprocess import sys@@ -51,6 +53,12 @@ DIRECTORY_FILE = "DIRECTORY.md" +# How many `gh pr view` calls to run concurrently. The slow "first pass" is one+# network round-trip per open PR, so it is I/O-bound and gains a lot from+# concurrency; the cap keeps us polite to the GitHub API and avoids secondary+# rate limits. Override with the PR_FILE_MAP_CONCURRENCY environment variable.+DEFAULT_CONCURRENCY = 10+ # Open PRs to skip in the report, e.g. [123, 456, 789] ignores #123, #456, #789. ignore_pull_request: set[int] = {15105, 15142, 15356} @@ -118,12 +126,67 @@ def get_open_prs() -> list[dict]: return [pr for pr in json.loads(raw) if pr["number"] not in ignore] -def get_pr_files(pr_number: int) -> list[str]:- raw = run_gh(["pr", "view", str(pr_number), "--json", "files"])+async def run_gh_async(args: list[str], semaphore: asyncio.Semaphore) -> str:+ """Async counterpart of run_gh, throttled by a shared semaphore.++ The semaphore bounds how many `gh` subprocesses run at once so we speed up+ the many-round-trip "first pass" without flooding the GitHub API.+ """+ async with semaphore:+ try:+ proc = await asyncio.create_subprocess_exec(+ "gh",+ *args,+ stdout=asyncio.subprocess.PIPE,+ stderr=asyncio.subprocess.PIPE,+ )+ except FileNotFoundError:+ sys.exit("Error: 'gh' (GitHub CLI) is not installed or not in PATH.")+ stdout, stderr = await proc.communicate()+ if proc.returncode != 0:+ sys.exit(f"Error running 'gh {' '.join(args)}':\n{stderr.decode().strip()}")+ return stdout.decode()+++async def get_pr_files_async(pr_number: int, semaphore: asyncio.Semaphore) -> list[str]:+ raw = await run_gh_async(+ ["pr", "view", str(pr_number), "--json", "files"], semaphore+ ) data = json.loads(raw) return [f["path"] for f in data.get("files", [])] +async def gather_pr_files(+ pr_numbers: list[int], concurrency: int+) -> dict[int, list[str]]:+ """Fetch each PR's file list concurrently, capped at ``concurrency``.++ Returns a ``{pr_number: [paths]}`` mapping keyed in the same order as+ ``pr_numbers`` so downstream output stays deterministic.+ """+ semaphore = asyncio.Semaphore(concurrency)+ results = await asyncio.gather(+ *(get_pr_files_async(number, semaphore) for number in pr_numbers)+ )+ return dict(zip(pr_numbers, results))+++def resolve_concurrency() -> int:+ """Read PR_FILE_MAP_CONCURRENCY (a positive int) or fall back to the default."""+ raw = os.environ.get("PR_FILE_MAP_CONCURRENCY")+ if raw is None:+ return DEFAULT_CONCURRENCY+ try:+ value = int(raw)+ except ValueError:+ value = 0+ if value < 1:+ sys.exit(+ f"Error: PR_FILE_MAP_CONCURRENCY must be a positive integer, got {raw!r}."+ )+ return value++ def split_directory_conflicts( directory_prs: list[int], pr_to_files: dict[int, list[str]],@@ -209,19 +272,27 @@ def main() -> None: print(f"PR count from get_open_prs(): {pr_count}", file=sys.stderr) file_to_prs: dict[str, list[int]] = defaultdict(list)- pr_to_files: dict[int, list[str]] = {} touch_count = 0 # every (PR, file) pair; a file may be touched by many PRs - for pr in prs:- pr_number = pr["number"]- pr_files = get_pr_files(pr_number)- pr_to_files[pr_number] = pr_files+ # First pass: one `gh pr view` per PR. This is the slow, network-bound part,+ # so fetch them concurrently (bounded by resolve_concurrency()).+ concurrency = resolve_concurrency()+ pr_numbers = [pr["number"] for pr in prs]+ print(+ f"Fetching files for {pr_count} PRs "+ f"(up to {concurrency} concurrent gh calls)...",+ file=sys.stderr,+ )+ pr_to_files = asyncio.run(gather_pr_files(pr_numbers, concurrency))++ for pr_number in pr_numbers:+ pr_files = pr_to_files[pr_number] touch_count += len(pr_files) for path in pr_files: file_to_prs[path].append(pr_number) distinct_count = len(file_to_prs) print(- f"File touches from get_pr_files(): {touch_count} "+ f"File touches from get_pr_files_async(): {touch_count} " f"across {distinct_count} distinct files", file=sys.stderr, )