TheRock/build_tools/github_actions/manifest_utils.py

#!/usr/bin/env python3
"""
Helpers for manifest generation.
"""

from dataclasses import dataclass
from pathlib import Path
import shlex
import subprocess
import sys


def log(*args, **kwargs) -> None:
    """Consistent logging helper for manifest generation scripts."""
    print(*args, **kwargs)
    sys.stdout.flush()


def warn(*args, **kwargs) -> None:
    """Consistent warning helper for manifest generation scripts."""
    print("WARNING:", *args, file=sys.stderr, **kwargs)
    sys.stderr.flush()


@dataclass(frozen=True)
class GitSourceInfo:
    """Git source info for a repository checkout."""

    repo: str
    commit: str
    branch: str | None = None
    version: str | None = None

    def to_dict(self) -> dict[str, str]:
        d = {"repo": self.repo, "commit": self.commit}
        if self.branch is not None:
            d["branch"] = self.branch
        if self.version is not None:
            d["version"] = self.version
        return d


def capture(args: list[str | Path], cwd: Path) -> str:
    args = [str(arg) for arg in args]
    log(f"++ Exec [{cwd}]$ {shlex.join(args)}")
    return (
        subprocess.check_output(
            args,
            cwd=str(cwd),
            stdin=subprocess.DEVNULL,
        )
        .decode()
        .strip()
    )


def capture_optional(args: list[str | Path], cwd: Path) -> str | None:
    """Like capture(), but returns None on failure."""
    args = [str(arg) for arg in args]
    log(f"++ Exec [{cwd}]$ {shlex.join(args)}")
    try:
        out = (
            subprocess.check_output(
                args,
                cwd=str(cwd),
                stdin=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
            .decode()
            .strip()
        )
    except (subprocess.CalledProcessError, FileNotFoundError):
        return None
    return out or None


def git_upstream_commit(dirpath: Path, *, label: str) -> str:
    """Return the true upstream commit, skipping any local commit.

    The TheRock checkout wrapper (repo_management.py) tags the real upstream
    commit as THEROCK_UPSTREAM_DIFFBASE before HIPIFY runs. We prefer
    that tag; if it is absent (plain clone) we fall back to HEAD.
    """
    commit = capture_optional(
        ["git", "rev-parse", "THEROCK_UPSTREAM_DIFFBASE"], cwd=dirpath
    )
    if commit:
        log(f"{label} upstream commit (via THEROCK_UPSTREAM_DIFFBASE): {commit}")
    else:
        commit = capture(["git", "rev-parse", "HEAD"], cwd=dirpath)
        log(f"{label} upstream commit (HEAD fallback): {commit}")
    return commit


def git_head(dirpath: Path, *, label: str) -> GitSourceInfo:
    """Return commit + origin repo for a git checkout."""
    dirpath = dirpath.resolve()

    if not dirpath.exists():
        raise FileNotFoundError(
            f"{label}: directory does not exist: {dirpath}\n"
            "This indicates a misconfigured workflow or incomplete checkout."
        )

    if not (dirpath / ".git").exists():
        raise FileNotFoundError(
            f"{label}: not a git checkout (missing .git): {dirpath}\n"
            "Manifest generation requires git commit hash and origin repo."
        )

    commit = git_upstream_commit(dirpath, label=label)
    repo = capture(["git", "remote", "get-url", "origin"], cwd=dirpath)
    return GitSourceInfo(commit=commit, repo=repo)


def git_branch_best_effort(dirpath: Path) -> str | None:
    """Return current branch name if on a real branch; None if detached/unknown."""
    dirpath = dirpath.resolve()

    # Most reliable when on a branch; fails in detached HEAD.
    b = capture_optional(
        ["git", "symbolic-ref", "--quiet", "--short", "HEAD"], cwd=dirpath
    )
    if b and b != "HEAD":
        return b

    # Fallback. Returns empty on detached.
    b = capture_optional(["git", "branch", "--show-current"], cwd=dirpath)
    if b and b != "HEAD":
        return b

    return None


def resolve_branch(*, inferred: str | None, provided: str | None) -> str | None:
    """Choose inferred branch if available; else provided; else None."""
    if inferred:
        return inferred
    if provided:
        return provided
    return None


def detect_therock_source_info(repo_root: Path) -> GitSourceInfo:
    """Detect TheRock commit, repo, and branch from a local git checkout.

    Falls back to "unknown" for fields that cannot be determined (e.g. when
    not inside a git worktree).
    """
    commit = capture_optional(["git", "rev-parse", "HEAD"], cwd=repo_root)
    if commit is None:
        warn(f"Could not detect TheRock commit from {repo_root}; using 'unknown'")
        commit = "unknown"

    repo_url = capture_optional(["git", "remote", "get-url", "origin"], cwd=repo_root)
    if repo_url is None:
        warn(
            f"Could not detect TheRock origin from {repo_root}; "
            "using https://github.com/ROCm/TheRock"
        )
        repo_url = "https://github.com/ROCm/TheRock"
    repo_url = repo_url.removesuffix(".git")
    branch = git_branch_best_effort(repo_root)
    if branch is None:
        warn(f"Could not detect TheRock branch from {repo_root}; using 'unknown'")
        branch = "unknown"
    return GitSourceInfo(commit=commit, repo=repo_url, branch=branch)


def normalize_python_version_for_filename(python_version: str) -> str:
    """Normalize python version strings for filenames.

    Examples:
      "py3.12" -> "3.12"
      "3.12"   -> "3.12"
    """
    py = python_version.strip()
    if py.startswith("py"):
        py = py[2:]
    return py


def normalize_ref_for_filename(ref: str) -> str:
    """Normalize a git ref for filenames by replacing '/' with '-'.

    Examples:
      "nightly"                -> "nightly"
      "release/0.4.28"         -> "release-0.4.28"
      "users/alice/experiment" -> "users-alice-experiment"
    """
    return ref.replace("/", "-")