"""Render UML/BPMN/C4 fenced blocks to PNG (PlantUML jar / Kroki local / remote).""" from __future__ import annotations import hashlib import logging import os import shutil import subprocess import tempfile from dataclasses import dataclass from pathlib import Path from typing import Literal import requests from . import package_dir DIAGRAM_LANGS = frozenset({"uml", "plantuml", "bpmn", "c4"}) FallbackMode = Literal["local", "remote", "off"] DEFAULT_KROKI_URL = "http://localhost:8000" REMOTE_KROKI_URL = "https://kroki.io" _log = logging.getLogger(__name__) @dataclass class DiagramConfig: plantuml_jar: str | None = None kroki_url: str | None = None fallback: FallbackMode = "remote" cache_dir: str | None = None _CONFIG = DiagramConfig() def configure_diagrams( plantuml_jar: str | None = None, kroki_url: str | None = None, fallback: FallbackMode = "remote", cache_dir: str | None = None, ) -> None: global _CONFIG _CONFIG = DiagramConfig( plantuml_jar=plantuml_jar or os.environ.get("PLANTUML_JAR"), kroki_url=kroki_url or os.environ.get("KROKI_URL"), fallback=fallback, cache_dir=cache_dir, ) def diagrams_dir() -> Path: return Path(package_dir()) / "diagrams" def prepare_source(lang: str, source: str) -> tuple[str, str]: """Return (prepared_source, kroki_diagram_type).""" lang = (lang or "uml").lower().strip() text = source.strip() if lang == "bpmn": if "@startbpmn" not in text.lower() and "\n" ) body = text if body.lower().startswith("@startuml"): lines = body.splitlines() text = lines[0] + "\n" + includes + "\n".join(lines[1:]) else: text = f"@startuml\n{includes}{body}\n@enduml" elif not text.lower().startswith("@start"): text = f"@startuml\n{text}\n@enduml" return text, "plantuml" # uml / plantuml if not text.lower().startswith("@start"): text = f"@startuml\n{text}\n@enduml" return text, "plantuml" def _cache_path(source: str, cache_dir: str | None = None) -> Path: root = cache_dir or _CONFIG.cache_dir if not root: wd = os.environ.get("WORKING_DIR", ".") root = os.path.join(wd, ".md2gost-cache") Path(root).mkdir(parents=True, exist_ok=True) digest = hashlib.sha256(source.encode("utf-8")).hexdigest()[:24] return Path(root) / f"{digest}.png" def _java_available() -> bool: return shutil.which("java") is not None def _render_plantuml_jar(source: str, out_png: Path, jar: str) -> bool: if not jar or not os.path.isfile(jar) or not _java_available(): return False with tempfile.TemporaryDirectory() as tmp: src = Path(tmp) / "diagram.puml" src.write_text(source, encoding="utf-8") cmd = [ "java", "-jar", jar, "-tpng", "-charset", "UTF-8", f"-I{diagrams_dir()}", "-o", tmp, str(src), ] try: proc = subprocess.run( cmd, capture_output=True, text=True, timeout=120, check=False, ) except (OSError, subprocess.TimeoutExpired) as e: _log.warning("PlantUML jar failed: %s", e) return False produced = Path(tmp) / "diagram.png" if proc.returncode != 0 or not produced.is_file(): _log.warning("PlantUML jar error: %s", proc.stderr or proc.stdout) return False shutil.copyfile(produced, out_png) return True def _render_kroki(source: str, diagram_type: str, base_url: str, out_png: Path) -> bool: url = base_url.rstrip("/") + f"/{diagram_type}/png" try: resp = requests.post( url, data=source.encode("utf-8"), headers={"Content-Type": "text/plain"}, timeout=60, ) if resp.status_code != 200 or not resp.content.startswith(b"\x89PNG"): _log.warning("Kroki %s → HTTP %s", url, resp.status_code) return False out_png.write_bytes(resp.content) return True except requests.RequestException as e: _log.warning("Kroki request failed (%s): %s", url, e) return False def render_diagram( lang: str, source: str, *, plantuml_jar: str | None = None, kroki_url: str | None = None, fallback: FallbackMode | None = None, cache_dir: str | None = None, ) -> str: """Render diagram to PNG path. Raises RuntimeError if all backends fail.""" prepared, kroki_type = prepare_source(lang, source) out = _cache_path(prepared, cache_dir) if out.is_file() and out.stat().st_size > 0: return str(out) jar = plantuml_jar if plantuml_jar is not None else _CONFIG.plantuml_jar jar = jar or os.environ.get("PLANTUML_JAR") local_kroki = kroki_url if kroki_url is not None else _CONFIG.kroki_url local_kroki = local_kroki or os.environ.get("KROKI_URL") or DEFAULT_KROKI_URL mode: FallbackMode = fallback if fallback is not None else _CONFIG.fallback # 1) PlantUML jar if _render_plantuml_jar(prepared, out, jar or ""): return str(out) # 2) Local Kroki if _render_kroki(prepared, kroki_type, local_kroki, out): return str(out) if mode == "off": raise RuntimeError( "Не удалось отрендерить диаграмму локально " "(задайте --plantuml-jar или KROKI_URL; remote fallback отключён)" ) if mode == "local": raise RuntimeError( "Локальный рендер диаграммы недоступен " "(Java+plantuml.jar или локальный Kroki)" ) # 3) Remote fallback _log.warning( "Диаграмма: локальный рендер недоступен, использую remote %s", REMOTE_KROKI_URL, ) if _render_kroki(prepared, kroki_type, REMOTE_KROKI_URL, out): return str(out) raise RuntimeError("Не удалось отрендерить диаграмму (PlantUML/Kroki)")