iron/scripts/generate_readme_demo.py

from __future__ import annotations

import shutil
import sys
from pathlib import Path

from PIL import Image, ImageOps
from playwright.sync_api import sync_playwright

ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from scripts.ui_playwright_smoke import (
    OUTPUT_DIR,
    create_temp_preview_file,
    create_temp_upload_file,
    ensure_frontend_built,
    find_free_port,
    prepare_database,
    start_server,
    wait_for_server,
)

DEMO_DIR = ROOT / "output" / "readme-demo"
RUNTIME_DIR = DEMO_DIR / "runtime"
ASSET_DIR = ROOT / "docs" / "assets"
GIF_PATH = ASSET_DIR / "iron-demo.gif"
FRAME_SIZE = (1200, 760)
FRAME_DURATION_MS = 1300


def main() -> int:
    DEMO_DIR.mkdir(parents=True, exist_ok=True)
    ASSET_DIR.mkdir(parents=True, exist_ok=True)
    if RUNTIME_DIR.exists():
        shutil.rmtree(RUNTIME_DIR)
    RUNTIME_DIR.mkdir(parents=True, exist_ok=True)

    ensure_frontend_built()

    port = find_free_port()
    base_url = f"http://127.0.0.1:{port}"
    runtime_db = RUNTIME_DIR / "readme-demo.db"
    prepare_database(runtime_db)
    process = start_server(port, runtime_db)

    try:
        wait_for_server(base_url)
        frame_paths = capture_demo_frames(base_url)
        build_gif(frame_paths, GIF_PATH)
    finally:
        process.terminate()
        try:
            process.wait(timeout=10)
        except Exception:
            process.kill()

    print(f"README demo GIF created at {GIF_PATH}")
    return 0


def capture_demo_frames(base_url: str) -> list[Path]:
    frame_paths: list[Path] = []
    with sync_playwright() as playwright:
        browser = playwright.chromium.launch(headless=True)
        context = browser.new_context(viewport={"width": 1440, "height": 1024}, device_scale_factor=1)
        page = context.new_page()

        page.goto(f"{base_url}/app/login", wait_until="networkidle")
        frame_paths.append(capture_frame(page, "01-login"))

        page.get_by_label("Username").fill("admin")
        page.get_by_label("Password").fill("changeme-iron")
        page.get_by_role("button", name="Sign in").click()
        page.wait_for_url(f"{base_url}/app/files")
        page.locator(".files-header").filter(has_text="Personal").first.wait_for()
        frame_paths.append(capture_frame(page, "02-files"))

        page.locator(".panel-toolbar").get_by_role("button", name="New").click()
        page.get_by_role("dialog").get_by_label("Name").fill("demo-folder")
        page.get_by_role("dialog").get_by_role("button", name="Confirm").click()
        page.locator(".drive-item strong", has_text="demo-folder").first.wait_for()
        frame_paths.append(capture_frame(page, "03-folder"))

        temp_upload = create_temp_upload_file()
        preview_upload = create_temp_preview_file()
        try:
            page.locator(".files-command-bar").get_by_role("button", name="Upload", exact=True).click()
            page.locator("input[type='file']").set_input_files(str(temp_upload))
            page.locator(".drive-item strong", has_text="playwright-upload.txt").first.wait_for(timeout=15000)

            page.locator(".files-command-bar").get_by_role("button", name="Upload", exact=True).click()
            page.locator("input[type='file']").set_input_files(str(preview_upload))
            page.locator(".drive-item strong", has_text="playwright-preview.png").first.wait_for(timeout=15000)
        finally:
            temp_upload.unlink(missing_ok=True)
            preview_upload.unlink(missing_ok=True)
        frame_paths.append(capture_frame(page, "04-upload"))

        page.locator(".drive-item strong", has_text="playwright-preview.png").first.click()
        page.locator("img.preview-frame").wait_for(timeout=15000)
        frame_paths.append(capture_frame(page, "05-preview"))

        page.locator(".detail-panel").get_by_role("button", name="More Info").click()
        page.locator(".detail-drawer", has_text="Replica Info").wait_for()
        page.locator(".detail-drawer", has_text="Desired Backends").wait_for()
        frame_paths.append(capture_frame(page, "06-storage"))

        page.goto(f"{base_url}/app/storage", wait_until="networkidle")
        frame_paths.append(capture_frame(page, "07-backends"))

        page.goto(f"{base_url}/app/policies", wait_until="networkidle")
        page.get_by_role("button", name="Save Policy").wait_for()
        frame_paths.append(capture_frame(page, "08-policies"))

        page.goto(f"{base_url}/app/jobs", wait_until="networkidle")
        frame_paths.append(capture_frame(page, "09-jobs"))

        context.close()
        browser.close()
    return frame_paths


def capture_frame(page, name: str) -> Path:  # type: ignore[no-untyped-def]
    path = DEMO_DIR / f"{name}.png"
    page.screenshot(path=str(path), full_page=True)
    return path


def build_gif(frame_paths: list[Path], output_path: Path) -> None:
    frames = []
    for frame_path in frame_paths:
        with Image.open(frame_path) as image:
            framed = ImageOps.contain(image.convert("RGB"), FRAME_SIZE)
            canvas = Image.new("RGB", FRAME_SIZE, color=(245, 247, 250))
            offset = ((FRAME_SIZE[0] - framed.width) // 2, (FRAME_SIZE[1] - framed.height) // 2)
            canvas.paste(framed, offset)
            frames.append(canvas)

    if not frames:
        raise RuntimeError("No frames captured for README demo GIF.")

    frames[0].save(
        output_path,
        save_all=True,
        append_images=frames[1:],
        duration=FRAME_DURATION_MS,
        loop=0,
        optimize=True,
        disposal=2,
    )


if __name__ == "__main__":
    raise SystemExit(main())