Zivro/Report/render_uml_png.py

#!/usr/bin/env python3
"""
Render PlantUML diagrams (.puml) to high-resolution PNG files via Kroki.

Why this script:
- local `plantuml` may require X11 in some environments;
- this script uses Kroki HTTP API and can raise output quality via `skinparam dpi`.

Examples:
    python3 Report/render_uml_png.py
    python3 Report/render_uml_png.py --dpi 360 --glob "*.puml"
    python3 Report/render_uml_png.py --input-dir Report/uml --timeout 120
"""

from __future__ import annotations

import argparse
import io
from pathlib import Path
import re
from typing import Iterable

import requests
from PIL import Image


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Render .puml diagrams to PNG with configurable DPI."
    )
    parser.add_argument(
        "--input-dir",
        default="Report/uml",
        help="Directory with .puml files (default: Report/uml)",
    )
    parser.add_argument(
        "--glob",
        default="*.puml",
        help="Glob pattern inside input-dir (default: *.puml)",
    )
    parser.add_argument(
        "--dpi",
        type=int,
        default=300,
        help="Output DPI via PlantUML skinparam (default: 300)",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=120,
        help="HTTP timeout in seconds (default: 120)",
    )
    parser.add_argument(
        "--kroki-url",
        default="https://kroki.io/plantuml/png",
        help="Kroki PlantUML PNG endpoint",
    )
    parser.add_argument(
        "--min-dpi",
        type=int,
        default=96,
        help="Lower bound for auto DPI fallback (default: 96)",
    )
    parser.add_argument(
        "--max-dim",
        type=int,
        default=3900,
        help=(
            "Target max PNG dimension during auto-fit. "
            "Keep below backend hard limits (default: 3900)"
        ),
    )
    return parser.parse_args()


def iter_puml_files(input_dir: Path, pattern: str) -> Iterable[Path]:
    return sorted(p for p in input_dir.glob(pattern) if p.is_file())


def inject_or_replace_skinparam(plantuml_text: str, key: str, value: int) -> str:
    """
    Ensure diagram has the requested `skinparam <key>`.
    - If skinparam exists, replace it.
    - Otherwise insert after @startuml line.
    """
    line_value = f"skinparam {key} {value}"
    pattern = rf"(?im)^\s*skinparam\s+{re.escape(key)}\s+\d+\s*$"

    if re.search(pattern, plantuml_text):
        return re.sub(pattern, line_value, plantuml_text)

    lines = plantuml_text.splitlines()
    for i, line in enumerate(lines):
        if line.strip().lower().startswith("@startuml"):
            lines.insert(i + 1, line_value)
            return "\n".join(lines) + ("\n" if plantuml_text.endswith("\n") else "")

    # Fallback: if @startuml is missing, prepend anyway.
    return line_value + "\n" + plantuml_text


def render_one(
    puml_file: Path,
    kroki_url: str,
    dpi: int,
    min_dpi: int,
    max_dim: int,
    timeout: int,
) -> Path:
    source = puml_file.read_text(encoding="utf-8")
    source = inject_or_replace_skinparam(source, "padding", 24)

    current_dpi = dpi
    png_bytes: bytes | None = None

    while True:
        payload = inject_or_replace_skinparam(source, "dpi", current_dpi)
        response = requests.post(
            kroki_url,
            data=payload.encode("utf-8"),
            timeout=timeout,
        )
        response.raise_for_status()
        candidate = response.content

        with Image.open(io.BytesIO(candidate)) as im:
            width, height = im.size

        # Kroki/PlantUML PNG responses may hit hard canvas limits.
        # If we are close to that ceiling, lower DPI and retry.
        if (width > max_dim or height > max_dim) and current_dpi > min_dpi:
            next_dpi = max(min_dpi, int(current_dpi * 0.82))
            if next_dpi == current_dpi:
                next_dpi = current_dpi - 1
            current_dpi = max(min_dpi, next_dpi)
            continue

        png_bytes = candidate
        break

    if png_bytes is None:
        raise RuntimeError(f"Failed to render diagram: {puml_file}")

    out_file = puml_file.with_suffix(".png")
    out_file.write_bytes(png_bytes)
    return out_file


def main() -> int:
    args = parse_args()

    input_dir = Path(args.input_dir).resolve()
    if not input_dir.exists():
        raise FileNotFoundError(f"Input directory not found: {input_dir}")

    files = list(iter_puml_files(input_dir, args.glob))
    if not files:
        print(f"No files matched in {input_dir} with pattern {args.glob}")
        return 0

    print(f"Rendering {len(files)} diagrams from: {input_dir}")
    print(f"DPI: {args.dpi}")

    success = 0
    for puml in files:
        out = render_one(
            puml_file=puml,
            kroki_url=args.kroki_url,
            dpi=args.dpi,
            min_dpi=args.min_dpi,
            max_dim=args.max_dim,
            timeout=args.timeout,
        )
        print(f"created {out}")
        success += 1

    print(f"Done: {success}/{len(files)} PNG files generated.")
    return 0


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