TheAlgorithms/Python · #15374

pr_file_map.py: fetch PR files concurrently with asyncio

priya-sundaram-dev · merged Sep 18, 20261 files · 80 + / 9
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,     )