Igor20264
b38661f588
Python application / build (push) Has been cancelled
- Add\Rework UI - Add Split Table and Listing - Add Support Customazeble schems
496 lines
16 KiB
Python
496 lines
16 KiB
Python
"""Render UML / Mermaid / scheme fenced blocks to PNG (+ optional SVG)."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import logging
|
||
import os
|
||
import re
|
||
import shutil
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Literal
|
||
|
||
import requests
|
||
|
||
from . import package_dir
|
||
from .diagram_schemes import (
|
||
configure_schemes,
|
||
ensure_user_schemes,
|
||
is_diagram_lang,
|
||
prepare_with_schemes,
|
||
)
|
||
|
||
# Back-compat: static set used by older tests; runtime uses is_diagram_lang().
|
||
DIAGRAM_LANGS = frozenset(
|
||
{
|
||
"uml",
|
||
"plantuml",
|
||
"c4",
|
||
"c4context",
|
||
"c4component",
|
||
"usecase",
|
||
"bpmn",
|
||
"mermaid",
|
||
"mmd",
|
||
}
|
||
)
|
||
FallbackMode = Literal["local", "remote", "off"]
|
||
DiagramFormat = Literal["png", "svg"]
|
||
|
||
DEFAULT_KROKI_URL = "http://localhost:8000"
|
||
REMOTE_KROKI_URL = "https://kroki.io"
|
||
|
||
# Sharper raster when Word shrinks the figure to page width (display size stays ×1).
|
||
DEFAULT_DIAGRAM_SCALE = 2.0
|
||
# Avoid PlantUML clipping large scaled diagrams (default limit is 4096).
|
||
PLANTUML_LIMIT_SIZE = "8192"
|
||
|
||
# Pinned official jar (GPL). Bundled at exe build; otherwise downloaded once to user cache.
|
||
PLANTUML_VERSION = "1.2025.4"
|
||
PLANTUML_JAR_URL = (
|
||
f"https://github.com/plantuml/plantuml/releases/download/"
|
||
f"v{PLANTUML_VERSION}/plantuml-{PLANTUML_VERSION}.jar"
|
||
)
|
||
|
||
_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
|
||
schemes_path: str | None = None
|
||
diagram_format: DiagramFormat = "png"
|
||
diagram_scale: float = DEFAULT_DIAGRAM_SCALE
|
||
|
||
|
||
@dataclass
|
||
class DiagramRenderResult:
|
||
png_path: str
|
||
svg_path: str | None = None
|
||
# PlantUML render scale; Image divides display size by this to keep ×1 layout.
|
||
pixel_scale: float = 1.0
|
||
|
||
|
||
_CONFIG = DiagramConfig()
|
||
|
||
|
||
def configure_diagrams(
|
||
plantuml_jar: str | None = None,
|
||
kroki_url: str | None = None,
|
||
fallback: FallbackMode = "remote",
|
||
cache_dir: str | None = None,
|
||
schemes_path: str | None = None,
|
||
md_dir: str | None = None,
|
||
diagram_format: DiagramFormat = "png",
|
||
diagram_scale: float | None = None,
|
||
) -> None:
|
||
global _CONFIG
|
||
ensure_user_schemes()
|
||
configure_schemes(schemes_path=schemes_path, md_dir=md_dir)
|
||
fmt: DiagramFormat = "svg" if diagram_format == "svg" else "png"
|
||
scale = DEFAULT_DIAGRAM_SCALE if diagram_scale is None else float(diagram_scale)
|
||
if scale < 1:
|
||
scale = 1.0
|
||
_CONFIG = DiagramConfig(
|
||
plantuml_jar=resolve_plantuml_jar(plantuml_jar or None, download=False),
|
||
kroki_url=kroki_url or os.environ.get("KROKI_URL"),
|
||
fallback=fallback,
|
||
cache_dir=cache_dir,
|
||
schemes_path=schemes_path,
|
||
diagram_format=fmt,
|
||
diagram_scale=scale,
|
||
)
|
||
|
||
|
||
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)."""
|
||
return prepare_with_schemes(lang, source)
|
||
|
||
|
||
def _cache_root(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 = Path(root)
|
||
path.mkdir(parents=True, exist_ok=True)
|
||
return path
|
||
|
||
|
||
def _cache_digest(source: str) -> str:
|
||
return hashlib.sha256(source.encode("utf-8")).hexdigest()[:24]
|
||
|
||
|
||
def _cache_paths(source: str, cache_dir: str | None = None) -> tuple[Path, Path]:
|
||
root = _cache_root(cache_dir)
|
||
digest = _cache_digest(source)
|
||
return root / f"{digest}.png", root / f"{digest}.svg"
|
||
|
||
|
||
def _java_available() -> bool:
|
||
return shutil.which("java") is not None
|
||
|
||
|
||
def vendor_plantuml_path() -> Path:
|
||
return Path(package_dir()) / "vendor" / "plantuml.jar"
|
||
|
||
|
||
def cached_plantuml_path() -> Path:
|
||
if os.name == "nt":
|
||
root = Path(os.environ.get("LOCALAPPDATA") or Path.home()) / "md2gost"
|
||
else:
|
||
root = Path.home() / ".md2gost"
|
||
return root / "plantuml.jar"
|
||
|
||
|
||
def iter_plantuml_candidates(explicit: str | None = None) -> list[Path]:
|
||
paths: list[Path] = []
|
||
if explicit:
|
||
paths.append(Path(explicit))
|
||
env = os.environ.get("PLANTUML_JAR")
|
||
if env:
|
||
paths.append(Path(env))
|
||
paths.append(vendor_plantuml_path())
|
||
if getattr(sys, "frozen", False):
|
||
mei = getattr(sys, "_MEIPASS", None)
|
||
if mei:
|
||
paths.append(Path(mei) / "md2gost" / "vendor" / "plantuml.jar")
|
||
paths.append(Path(mei) / "vendor" / "plantuml.jar")
|
||
paths.append(Path(sys.executable).resolve().parent / "plantuml.jar")
|
||
paths.append(cached_plantuml_path())
|
||
seen: set[str] = set()
|
||
out: list[Path] = []
|
||
for path in paths:
|
||
key = str(path)
|
||
if key in seen:
|
||
continue
|
||
seen.add(key)
|
||
out.append(path)
|
||
return out
|
||
|
||
|
||
def resolve_plantuml_jar(explicit: str | None = None, *, download: bool = False) -> str | None:
|
||
"""Find a plantuml.jar: explicit path, env, bundled vendor, user cache; optionally download."""
|
||
for path in iter_plantuml_candidates(explicit):
|
||
if path.is_file() and path.stat().st_size > 1000:
|
||
return str(path)
|
||
if download:
|
||
dest = cached_plantuml_path()
|
||
if fetch_plantuml_jar(dest):
|
||
return str(dest)
|
||
return None
|
||
|
||
|
||
def fetch_plantuml_jar(dest: Path, url: str = PLANTUML_JAR_URL) -> bool:
|
||
"""Download official plantuml.jar to dest. Returns True on success."""
|
||
dest.parent.mkdir(parents=True, exist_ok=True)
|
||
tmp = dest.with_suffix(".jar.part")
|
||
try:
|
||
resp = requests.get(url, timeout=120, stream=True)
|
||
resp.raise_for_status()
|
||
size = 0
|
||
with open(tmp, "wb") as fh:
|
||
for chunk in resp.iter_content(chunk_size=65536):
|
||
if chunk:
|
||
fh.write(chunk)
|
||
size += len(chunk)
|
||
if size < 1000:
|
||
tmp.unlink(missing_ok=True)
|
||
return False
|
||
tmp.replace(dest)
|
||
_log.info("PlantUML jar: %s", dest)
|
||
return True
|
||
except Exception as exc:
|
||
_log.warning("Не удалось скачать plantuml.jar: %s", exc)
|
||
tmp.unlink(missing_ok=True)
|
||
return False
|
||
|
||
|
||
def diagram_engine_status(explicit_jar: str | None = None) -> str:
|
||
"""Human-readable status for the GUI."""
|
||
java = shutil.which("java")
|
||
jar = resolve_plantuml_jar(explicit_jar, download=False)
|
||
if java:
|
||
java_line = f"Java: есть ({java})"
|
||
else:
|
||
java_line = (
|
||
"Java: не найдена. Локальный PlantUML не запустится — "
|
||
"картинки пойдут через интернет (kroki.io)."
|
||
)
|
||
if jar:
|
||
jar_line = f"PlantUML: {jar}"
|
||
else:
|
||
jar_line = (
|
||
"PlantUML: файла нет. Нажмите «Скачать PlantUML» (нужна Java) "
|
||
"или оставьте как есть — картинки нарисует интернет (kroki.io)."
|
||
)
|
||
return java_line + "\n" + jar_line
|
||
|
||
|
||
def _ensure_plantuml_png_scale(source: str, scale: float = DEFAULT_DIAGRAM_SCALE) -> str:
|
||
"""Inject `scale N` into PlantUML source unless the author already set scale/dpi."""
|
||
if scale <= 1:
|
||
return source
|
||
if re.search(r"(?im)^\s*scale\b", source):
|
||
return source
|
||
if re.search(r"(?im)^\s*skinparam\s+dpi\b", source):
|
||
return source
|
||
lines = source.splitlines()
|
||
if not lines:
|
||
return source
|
||
insert_at = 0
|
||
if lines[0].strip().lower().startswith("@start"):
|
||
insert_at = 1
|
||
scale_line = f"scale {scale:g}"
|
||
lines.insert(insert_at, scale_line)
|
||
return "\n".join(lines) + ("\n" if source.endswith("\n") else "")
|
||
|
||
|
||
def _is_png(data: bytes) -> bool:
|
||
return data.startswith(b"\x89PNG")
|
||
|
||
|
||
def _is_svg(data: bytes) -> bool:
|
||
head = data.lstrip()[:200].lower()
|
||
return head.startswith(b"<svg") or head.startswith(b"<?xml") or b"<svg" in head
|
||
|
||
|
||
def _render_plantuml_jar_one(
|
||
source: str,
|
||
out_path: Path,
|
||
jar: str,
|
||
fmt: Literal["png", "svg"],
|
||
) -> 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",
|
||
f"-Dplantuml.include.path={diagrams_dir().resolve()}",
|
||
f"-DPLANTUML_LIMIT_SIZE={PLANTUML_LIMIT_SIZE}",
|
||
"-jar", jar,
|
||
f"-t{fmt}",
|
||
"-charset", "UTF-8",
|
||
"-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) / f"diagram.{fmt}"
|
||
if proc.returncode != 0 or not produced.is_file():
|
||
_log.warning("PlantUML jar error: %s", proc.stderr or proc.stdout)
|
||
return False
|
||
data = produced.read_bytes()
|
||
if fmt == "png" and not _is_png(data):
|
||
return False
|
||
if fmt == "svg" and not _is_svg(data):
|
||
return False
|
||
shutil.copyfile(produced, out_path)
|
||
return True
|
||
|
||
|
||
def _render_plantuml_jar(
|
||
source: str,
|
||
out_png: Path,
|
||
jar: str,
|
||
*,
|
||
out_svg: Path | None = None,
|
||
) -> bool:
|
||
"""Render PNG; optionally also SVG from the same jar."""
|
||
if not _render_plantuml_jar_one(source, out_png, jar, "png"):
|
||
return False
|
||
if out_svg is not None:
|
||
if not _render_plantuml_jar_one(source, out_svg, jar, "svg"):
|
||
_log.warning("PlantUML jar: PNG ok, SVG failed")
|
||
return True
|
||
|
||
|
||
def _render_kroki(
|
||
source: str,
|
||
diagram_type: str,
|
||
base_url: str,
|
||
out_path: Path,
|
||
fmt: Literal["png", "svg"] = "png",
|
||
) -> bool:
|
||
url = base_url.rstrip("/") + f"/{diagram_type}/{fmt}"
|
||
try:
|
||
resp = requests.post(
|
||
url,
|
||
data=source.encode("utf-8"),
|
||
headers={"Content-Type": "text/plain"},
|
||
timeout=60,
|
||
)
|
||
if resp.status_code != 200:
|
||
_log.warning("Kroki %s → HTTP %s", url, resp.status_code)
|
||
return False
|
||
data = resp.content
|
||
if fmt == "png" and not _is_png(data):
|
||
_log.warning("Kroki %s → not PNG", url)
|
||
return False
|
||
if fmt == "svg" and not _is_svg(data):
|
||
_log.warning("Kroki %s → not SVG", url)
|
||
return False
|
||
out_path.write_bytes(data)
|
||
return True
|
||
except requests.RequestException as e:
|
||
_log.warning("Kroki request failed (%s): %s", url, e)
|
||
return False
|
||
|
||
|
||
def _render_kroki_pair(
|
||
source: str,
|
||
diagram_type: str,
|
||
base_url: str,
|
||
out_png: Path,
|
||
out_svg: Path | None,
|
||
) -> bool:
|
||
if not _render_kroki(source, diagram_type, base_url, out_png, "png"):
|
||
return False
|
||
if out_svg is not None:
|
||
if not _render_kroki(source, diagram_type, base_url, out_svg, "svg"):
|
||
_log.warning("Kroki: PNG ok, SVG failed (%s)", base_url)
|
||
return True
|
||
|
||
|
||
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,
|
||
diagram_format: DiagramFormat | None = None,
|
||
diagram_scale: float | None = None,
|
||
) -> DiagramRenderResult:
|
||
"""Render diagram to PNG (+ optional SVG). Raises RuntimeError if all backends fail."""
|
||
ensure_user_schemes()
|
||
prepared, kroki_type = prepare_source(lang, source)
|
||
scale = _CONFIG.diagram_scale if diagram_scale is None else float(diagram_scale)
|
||
if scale < 1:
|
||
scale = 1.0
|
||
# pixel_scale: only when we injected scale (author override → treat as 1 for display).
|
||
pixel_scale = 1.0
|
||
if kroki_type == "plantuml" and scale > 1:
|
||
before = prepared
|
||
prepared = _ensure_plantuml_png_scale(prepared, scale)
|
||
if prepared != before:
|
||
pixel_scale = scale
|
||
out_png, out_svg_path = _cache_paths(prepared, cache_dir)
|
||
fmt: DiagramFormat = (
|
||
diagram_format
|
||
if diagram_format is not None
|
||
else _CONFIG.diagram_format
|
||
)
|
||
want_svg = fmt == "svg"
|
||
svg_target = out_svg_path if want_svg else None
|
||
|
||
if out_png.is_file() and out_png.stat().st_size > 0:
|
||
has_svg = out_svg_path.is_file() and out_svg_path.stat().st_size > 0
|
||
if not want_svg or has_svg:
|
||
return DiagramRenderResult(
|
||
png_path=str(out_png),
|
||
svg_path=str(out_svg_path) if (want_svg and has_svg) else None,
|
||
pixel_scale=pixel_scale,
|
||
)
|
||
# PNG cached but SVG missing in svg mode — fall through to fill SVG.
|
||
|
||
jar = plantuml_jar if plantuml_jar is not None else _CONFIG.plantuml_jar
|
||
jar = resolve_plantuml_jar(jar or None, download=False)
|
||
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
|
||
|
||
def _result() -> DiagramRenderResult:
|
||
svg = None
|
||
if want_svg and out_svg_path.is_file() and out_svg_path.stat().st_size > 0:
|
||
svg = str(out_svg_path)
|
||
return DiagramRenderResult(
|
||
png_path=str(out_png), svg_path=svg, pixel_scale=pixel_scale,
|
||
)
|
||
|
||
# 1) PlantUML jar — only for plantuml type
|
||
if kroki_type == "plantuml":
|
||
if _render_plantuml_jar(prepared, out_png, jar or "", out_svg=svg_target):
|
||
return _result()
|
||
|
||
# 2) Local Kroki
|
||
if _render_kroki_pair(prepared, kroki_type, local_kroki, out_png, svg_target):
|
||
return _result()
|
||
|
||
def _png_fallback_or_raise(message: str) -> DiagramRenderResult:
|
||
if out_png.is_file() and out_png.stat().st_size > 0:
|
||
_log.warning("Диаграмма: SVG недоступен, вставляю только PNG")
|
||
return DiagramRenderResult(
|
||
png_path=str(out_png), svg_path=None, pixel_scale=pixel_scale,
|
||
)
|
||
raise RuntimeError(message)
|
||
|
||
if mode == "off":
|
||
if kroki_type == "plantuml":
|
||
return _png_fallback_or_raise(
|
||
"Не удалось отрендерить диаграмму локально "
|
||
"(задайте --plantuml-jar или KROKI_URL; remote fallback отключён)"
|
||
)
|
||
return _png_fallback_or_raise(
|
||
"Не удалось отрендерить диаграмму через локальный Kroki "
|
||
"(задайте --kroki-url / KROKI_URL; remote fallback отключён)"
|
||
)
|
||
if mode == "local":
|
||
if kroki_type == "plantuml":
|
||
return _png_fallback_or_raise(
|
||
"Локальный рендер диаграммы недоступен "
|
||
"(Java+plantuml.jar или локальный Kroki)"
|
||
)
|
||
return _png_fallback_or_raise(
|
||
"Локальный Kroki недоступен для диаграммы "
|
||
f"(тип {kroki_type}; задайте работающий --kroki-url)"
|
||
)
|
||
|
||
# 3) Remote fallback
|
||
_log.warning(
|
||
"Диаграмма: локальный рендер недоступен, использую remote %s",
|
||
REMOTE_KROKI_URL,
|
||
)
|
||
if _render_kroki_pair(prepared, kroki_type, REMOTE_KROKI_URL, out_png, svg_target):
|
||
return _result()
|
||
|
||
return _png_fallback_or_raise(
|
||
"Не удалось отрендерить диаграмму (PlantUML/Kroki)"
|
||
)
|
||
|
||
|
||
# Re-export for callers / tests
|
||
__all__ = [
|
||
"DIAGRAM_LANGS",
|
||
"DiagramRenderResult",
|
||
"configure_diagrams",
|
||
"diagram_engine_status",
|
||
"diagrams_dir",
|
||
"fetch_plantuml_jar",
|
||
"is_diagram_lang",
|
||
"prepare_source",
|
||
"render_diagram",
|
||
"resolve_plantuml_jar",
|
||
"PLANTUML_JAR_URL",
|
||
"PLANTUML_VERSION",
|
||
]
|