Add PR creation script.

tjprescott · tjprescott · commit 7c45b158c963 · 2026-05-29T12:46:01.000-07:00
diff --git a/scripts/create_api_review_pr.py b/scripts/create_api_review_pr.py
@@ -0,0 +1,371 @@
+#!/usr/bin/env python
+"""Create an API review PR for an Azure SDK Python package.
+
+Workflow:
+  1. Validate that ``--package-name`` exists under ``sdk/*/``.
+  2. Build the BASE branch (``base_{package}_{base_version}``):
+       - If ``--base`` is a tag (e.g. ``azure-ai-projects_1.0.0b1``): check out
+         the tag, generate API.md, then create the base branch off the latest
+         ``origin/main`` and commit the captured API.md onto it.
+       - If ``--base`` is omitted: create the base branch off ``origin/main``
+         and delete any existing API.md for the package (no-op if absent).
+  3. Build the REVIEW branch (``review_{package}_{target_version}``):
+       - If ``--target`` is omitted: use the latest ``origin/main``.
+       - Otherwise: check out the given branch.
+       Generate API.md on that ref, then commit it on a branch created off
+       the base branch.
+  4. Push both branches to ``origin`` and open a PR with title:
+       ``[API Review] {package} {target_version} (base {base_version})``
+
+Usage::
+
+    python scripts/create_api_review_pr.py --package-name azure-ai-projects
+    python scripts/create_api_review_pr.py --package-name azure-ai-projects \\
+        --base azure-ai-projects_1.0.0b1
+    python scripts/create_api_review_pr.py --package-name azure-ai-projects \\
+        --base azure-ai-projects_1.0.0b1 --target my-feature-branch
+
+Requires ``gh`` (GitHub CLI) authenticated against the repository, plus push
+access on the ``origin`` remote.
+"""
+
+from __future__ import annotations
+
+import argparse
+import glob
+import os
+import re
+import shutil
+import subprocess
+import sys
+import tempfile
+from typing import Optional
+
+
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+GENERATE_SCRIPT = os.path.join(REPO_ROOT, "scripts", "generate_api_text.py")
+EXPORT_SCRIPT = os.path.join(REPO_ROOT, "eng", "common", "scripts", "Export-APIViewMarkdown.ps1")
+REMOTE = "origin"
+MAIN_REF = f"{REMOTE}/main"
+
+
+# ---------------------------------------------------------------------------
+# Shell helpers
+# ---------------------------------------------------------------------------
+
+def run(cmd, *, cwd: str = REPO_ROOT, check: bool = True, capture: bool = False, env: Optional[dict] = None) -> subprocess.CompletedProcess:
+    """Run a command, echoing it first."""
+    printable = " ".join(cmd) if isinstance(cmd, list) else cmd
+    print(f"$ {printable}")
+    return subprocess.run(
+        cmd,
+        cwd=cwd,
+        check=check,
+        text=True,
+        capture_output=capture,
+        env=env,
+    )
+
+
+def git(*args: str, capture: bool = False, check: bool = True) -> subprocess.CompletedProcess:
+    return run(["git", *args], capture=capture, check=check)
+
+
+def git_out(*args: str) -> str:
+    return git(*args, capture=True).stdout.strip()
+
+
+# ---------------------------------------------------------------------------
+# Package + ref helpers
+# ---------------------------------------------------------------------------
+
+def find_package_dir(package_name: str) -> str:
+    """Locate ``sdk/*/{package_name}`` containing a pyproject.toml or setup.py."""
+    pattern = os.path.join(REPO_ROOT, "sdk", "*", package_name)
+    matches = [
+        m for m in glob.glob(pattern)
+        if os.path.isdir(m)
+        and (
+            os.path.exists(os.path.join(m, "pyproject.toml"))
+            or os.path.exists(os.path.join(m, "setup.py"))
+        )
+    ]
+    if not matches:
+        raise SystemExit(f"ERROR: package '{package_name}' not found under sdk/*/")
+    if len(matches) > 1:
+        raise SystemExit(f"ERROR: multiple matches for '{package_name}': {matches}")
+    return matches[0]
+
+
+def package_rel_dir(package_dir: str) -> str:
+    """Repo-relative POSIX path for the package directory."""
+    return os.path.relpath(package_dir, REPO_ROOT).replace(os.sep, "/")
+
+
+def api_md_path(package_dir: str) -> str:
+    return os.path.join(package_dir, "API.md")
+
+
+def api_md_rel(package_dir: str) -> str:
+    return f"{package_rel_dir(package_dir)}/API.md"
+
+
+_VERSION_RE = re.compile(r"""^\s*VERSION\s*[:=]\s*["']([^"']+)["']""", re.MULTILINE)
+
+
+def read_version(package_dir: str) -> str:
+    """Find and parse a ``_version.py`` (or ``version.py``) inside ``package_dir``."""
+    candidates = []
+    candidates.extend(glob.glob(os.path.join(package_dir, "**", "_version.py"), recursive=True))
+    candidates.extend(glob.glob(os.path.join(package_dir, "**", "version.py"), recursive=True))
+    for path in candidates:
+        try:
+            text = open(path, "r", encoding="utf-8").read()
+        except OSError:
+            continue
+        m = _VERSION_RE.search(text)
+        if m:
+            return m.group(1)
+    raise SystemExit(f"ERROR: could not find a version string in {package_dir}")
+
+
+def tag_exists(tag: str) -> bool:
+    result = git("rev-parse", "--verify", "--quiet", f"refs/tags/{tag}", capture=True, check=False)
+    return result.returncode == 0
+
+
+def ensure_clean_worktree() -> None:
+    status = git_out("status", "--porcelain")
+    if status:
+        raise SystemExit(
+            "ERROR: working tree is not clean. Commit or stash changes before running.\n"
+            + status
+        )
+
+
+def current_branch() -> str:
+    return git_out("rev-parse", "--abbrev-ref", "HEAD")
+
+
+def remote_branch_ref(branch: str) -> str:
+    """Return the ref name for a branch on ``REMOTE``, fetching it first."""
+    git("fetch", REMOTE, branch)
+    return f"{REMOTE}/{branch}"
+
+
+# ---------------------------------------------------------------------------
+# API.md generation
+# ---------------------------------------------------------------------------
+
+def generate_api_md(package_name: str, package_dir: str) -> bytes:
+    """Run ``generate_api_text.py`` for the package and return the bytes of the
+    resulting API.md. The file is also left on disk at its canonical location.
+    """
+    print(f"--- Generating API.md for {package_name} on {current_branch_or_sha()} ---")
+    run([sys.executable, GENERATE_SCRIPT, package_name])
+    path = api_md_path(package_dir)
+    if not os.path.exists(path):
+        raise SystemExit(f"ERROR: generate_api_text.py did not produce {path}")
+    with open(path, "rb") as f:
+        return f.read()
+
+
+def current_branch_or_sha() -> str:
+    name = git_out("rev-parse", "--abbrev-ref", "HEAD")
+    if name == "HEAD":
+        return git_out("rev-parse", "--short", "HEAD")
+    return name
+
+
+# ---------------------------------------------------------------------------
+# Main workflow
+# ---------------------------------------------------------------------------
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    p.add_argument("--package-name", required=True,
+                   help="Package directory name under sdk/*/ (e.g. azure-ai-projects)")
+    p.add_argument("--base", default=None,
+                   help="Tag to use as the API.md baseline, formatted as "
+                        "'{package-name}_{version}'. Omit to make the baseline empty.")
+    p.add_argument("--target", default=None,
+                   help="Branch containing the API to review. Omit to use the latest origin/main.")
+    return p.parse_args()
+
+
+def validate_base_tag(package_name: str, base: str) -> str:
+    """Validate the ``--base`` tag format/existence and return the version."""
+    if not base.startswith(f"{package_name}_"):
+        raise SystemExit(
+            f"ERROR: --base tag '{base}' must start with '{package_name}_'."
+        )
+    version = base[len(package_name) + 1:]
+    if not version:
+        raise SystemExit(f"ERROR: --base tag '{base}' is missing the version suffix.")
+    if not tag_exists(base):
+        raise SystemExit(f"ERROR: tag '{base}' does not exist in this repository.")
+    return version
+
+
+def write_bytes(path: str, data: bytes) -> None:
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    with open(path, "wb") as f:
+        f.write(data)
+
+
+def main() -> int:
+    args = parse_args()
+    package_name = args.package_name
+
+    package_dir = find_package_dir(package_name)
+    print(f"Found package at: {package_dir}")
+
+    ensure_clean_worktree()
+    original_branch = current_branch()
+    if original_branch == "HEAD":
+        raise SystemExit("ERROR: refusing to run from a detached HEAD.")
+
+    # Always fetch main once up-front.
+    git("fetch", REMOTE, "main")
+
+    # ---- Validate inputs --------------------------------------------------
+    base_version = "none"
+    if args.base is not None:
+        base_version = validate_base_tag(package_name, args.base)
+
+    target_ref: str
+    if args.target is None:
+        target_ref = MAIN_REF
+    else:
+        target_ref = remote_branch_ref(args.target)
+
+    # Cache the generate + export scripts (they may not exist on older refs we check out).
+    tmp_script_dir = tempfile.mkdtemp(prefix="apirev_script_")
+    cached_script = os.path.join(tmp_script_dir, "generate_api_text.py")
+    cached_export = os.path.join(tmp_script_dir, "Export-APIViewMarkdown.ps1")
+    shutil.copy2(GENERATE_SCRIPT, cached_script)
+    shutil.copy2(EXPORT_SCRIPT, cached_export)
+
+    try:
+        # ---- Step 1: capture base API.md content (if base is a tag) ------
+        base_api_bytes: Optional[bytes] = None
+        if args.base is not None:
+            print(f"\n=== Capturing baseline API.md from tag {args.base} ===")
+            git("checkout", "--detach", args.base)
+            base_api_bytes = _generate_with_cached_script(
+                cached_script, cached_export, package_name, package_dir
+            )
+
+        # ---- Step 2: capture target API.md content -----------------------
+        print(f"\n=== Capturing target API.md from {target_ref} ===")
+        git("checkout", "--detach", target_ref)
+        target_version = read_version(package_dir)
+        target_api_bytes = _generate_with_cached_script(
+            cached_script, cached_export, package_name, package_dir
+        )
+
+        # ---- Step 3: build base branch off origin/main -------------------
+        base_branch = f"base_{package_name}_{base_version}"
+        review_branch = f"review_{package_name}_{target_version}"
+
+        print(f"\n=== Creating base branch {base_branch} ===")
+        git("checkout", "-B", base_branch, MAIN_REF)
+
+        api_path = api_md_path(package_dir)
+        api_relative = api_md_rel(package_dir)
+
+        if base_api_bytes is not None:
+            write_bytes(api_path, base_api_bytes)
+            git("add", api_relative)
+            git("commit", "-m",
+                f"[API Review] Baseline API.md for {package_name} {base_version}")
+        else:
+            if os.path.exists(api_path):
+                git("rm", api_relative)
+                git("commit", "-m",
+                    f"[API Review] Remove API.md for {package_name} (empty baseline)")
+            else:
+                git("commit", "--allow-empty", "-m",
+                    f"[API Review] Empty baseline for {package_name}")
+
+        git("push", "--force-with-lease", REMOTE, base_branch)
+
+        # ---- Step 4: build review branch off base branch -----------------
+        print(f"\n=== Creating review branch {review_branch} ===")
+        git("checkout", "-B", review_branch, base_branch)
+        write_bytes(api_path, target_api_bytes)
+        git("add", api_relative)
+        # If the bytes happen to be identical to the base, commit empty so we
+        # still have something to PR.
+        diff = git("diff", "--cached", "--quiet", capture=True, check=False)
+        if diff.returncode == 0:
+            git("commit", "--allow-empty", "-m",
+                f"[API Review] API.md for {package_name} {target_version} (no diff vs baseline)")
+        else:
+            git("commit", "-m",
+                f"[API Review] API.md for {package_name} {target_version}")
+
+        git("push", "--force-with-lease", REMOTE, review_branch)
+
+        # ---- Step 5: open PR --------------------------------------------
+        title = f"[API Review] {package_name} {target_version} (base {base_version})"
+        body_lines = [
+            f"Automated API review PR for `{package_name}`.",
+            "",
+            f"- **Target:** `{args.target or 'origin/main'}` (version `{target_version}`)",
+            f"- **Baseline:** {'tag `' + args.base + '`' if args.base else '_empty_'} "
+            f"(version `{base_version}`)",
+            "",
+            "Generated by `scripts/create_api_review_pr.py`.",
+        ]
+        body = "\n".join(body_lines)
+
+        print(f"\n=== Opening PR ===")
+        compare_url = (
+            f"https://github.com/Azure/azure-sdk-for-python/compare/"
+            f"{base_branch}...{review_branch}?expand=1"
+        )
+        pr_result = run([
+            "gh", "pr", "create",
+            "--repo", "Azure/azure-sdk-for-python",
+            "--base", base_branch,
+            "--head", review_branch,
+            "--title", title,
+            "--body", body,
+            "--draft",
+        ], check=False)
+        if pr_result.returncode != 0:
+            print(
+                "\nWARNING: `gh pr create` failed. Both branches were pushed "
+                "successfully -- open the PR manually here:\n"
+                f"  {compare_url}\n"
+                f"  Title: {title}"
+            )
+
+        return 0
+
+    finally:
+        # Restore the user's original branch.
+        try:
+            git("checkout", original_branch, check=False)
+        finally:
+            shutil.rmtree(tmp_script_dir, ignore_errors=True)
+
+
+def _generate_with_cached_script(cached_script: str, cached_export: str, package_name: str, package_dir: str) -> bytes:
+    """Run the cached copy of generate_api_text.py against the currently
+    checked-out ref and return the bytes of the resulting API.md."""
+    print(f"--- Generating API.md on {current_branch_or_sha()} ---")
+    env = os.environ.copy()
+    env["AZSDK_REPO_ROOT"] = REPO_ROOT
+    env["AZSDK_EXPORT_SCRIPT"] = cached_export
+    run([sys.executable, cached_script, package_name], env=env)
+    path = api_md_path(package_dir)
+    if not os.path.exists(path):
+        raise SystemExit(f"ERROR: did not produce {path}")
+    with open(path, "rb") as f:
+        return f.read()
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/scripts/generate_api_text.py b/scripts/generate_api_text.py
@@ -14,10 +14,10 @@
 import tempfile
 
 
-REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+REPO_ROOT = os.environ.get("AZSDK_REPO_ROOT") or os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 APIVIEW_REQS = os.path.join(REPO_ROOT, "eng", "apiview_reqs.txt")
 AZURE_SDK_INDEX = "https://pkgs.dev.azure.com/azure-sdk/public/_packaging/azure-sdk-for-python/pypi/simple/"
-EXPORT_SCRIPT = os.path.join(REPO_ROOT, "eng", "common", "scripts", "Export-APIViewMarkdown.ps1")
+EXPORT_SCRIPT = os.environ.get("AZSDK_EXPORT_SCRIPT") or os.path.join(REPO_ROOT, "eng", "common", "scripts", "Export-APIViewMarkdown.ps1")
 
 
 def find_package_dir(package_name: str) -> str:
