Igor20264
@@ -0,0 +1,204 @@
|
||||
"""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 "<definitions" not in text.lower():
|
||||
# PlantUML BPMN dialect
|
||||
if not text.startswith("@start"):
|
||||
text = "@startbpmn\n" + text + "\n@endbpmn"
|
||||
return text, "bpmn"
|
||||
|
||||
if lang == "c4":
|
||||
if "!include" not in text and "!includeurl" not in text.lower():
|
||||
# Prefer PlantUML stdlib; also ship local stubs for jar -I path
|
||||
includes = (
|
||||
f"!include {diagrams_dir() / 'C4_Container.puml'}\n"
|
||||
if (diagrams_dir() / "C4_Container.puml").exists()
|
||||
else "!include <C4/C4_Container>\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)")
|
||||
Reference in New Issue
Block a user