scripts/apply_llama_pin.py

"""Merge a fragment printed by `record_llama_pin.py` into `llama_runtime.json`.

The recorder deliberately does not write the file: a digest belongs in the
repository only once somebody has looked at it. This does the mechanical half
of that -- read the fragment, put it in the right place, keep the file's
formatting and its explanatory comment intact -- and prints the resulting entry
so the change is still something a person reads before it is committed.

Usage::

    python scripts/record_llama_pin.py windows-x64-cuda > cuda-pin.json
    python scripts/apply_llama_pin.py cuda-pin.json

Refuses a fragment whose URLs are not GitHub, and refuses one whose digests are
not 64 hex characters, because a pin that fails those checks will be refused at
install time anyway and is better refused here where somebody is watching.
"""

from __future__ import annotations

import json
import re
import sys
from pathlib import Path

PIN_FILE = Path(__file__).resolve().parent.parent / "llama_runtime.json"
_SHA = re.compile(r"^[0-9a-f]{64}$")
_GITHUB = ("https://github.com/", "https://objects.githubusercontent.com/")


def _archives(entry: dict) -> list[dict]:
    listed = entry.get("archives")
    return list(listed) if isinstance(listed, list) else [entry]


def _check(key: str, entry: dict) -> list[str]:
    problems = []
    for archive in _archives(entry):
        url = str(archive.get("url", ""))
        digest = str(archive.get("sha256", "")).casefold()
        if not url.startswith(_GITHUB):
            problems.append(f"{key}: {url!r} is not a GitHub release URL")
        if not _SHA.match(digest):
            problems.append(f"{key}: {digest!r} is not a SHA-256 digest")
    if not _archives(entry):
        problems.append(f"{key}: names no archives")
    return problems


def main(argv: list[str]) -> int:
    if len(argv) != 1:
        print(__doc__, file=sys.stderr)
        return 2
    fragment = json.loads(Path(argv[0]).read_text(encoding="utf-8"))
    if not isinstance(fragment, dict) or not fragment:
        print("the fragment is not a JSON object of builds", file=sys.stderr)
        return 1

    problems = [line for key, entry in fragment.items() for line in _check(key, entry)]
    if problems:
        print("\n".join(problems), file=sys.stderr)
        return 1

    pins = json.loads(PIN_FILE.read_text(encoding="utf-8"))
    builds = pins.setdefault("builds", {})
    for key, entry in fragment.items():
        was = builds.get(key)
        builds[key] = entry
        print(f"{'replaced' if was else 'added'} {key}")
        print(json.dumps(entry, indent=2))
    PIN_FILE.write_text(json.dumps(pins, indent=2) + "\n", encoding="utf-8")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))