77 lines
2 KiB
Python
77 lines
2 KiB
Python
import base64
|
|
import zlib
|
|
from mkdocs.utils.cache import download_and_cache_url
|
|
from datetime import timedelta
|
|
from mkdocs.plugins import get_plugin_logger
|
|
from urllib.request import build_opener, install_opener, _opener
|
|
import hashlib
|
|
from functools import partial
|
|
from mkdocs.utils import write_file
|
|
import os
|
|
from py_svg_hush import filter_svg
|
|
from scour.scour import scourString
|
|
from slugify import slugify
|
|
|
|
logger = get_plugin_logger("kroki")
|
|
|
|
DIAGRAM_TYPES = {
|
|
"blockdiag",
|
|
"mermaid"
|
|
}
|
|
|
|
DIAGRAM_CACHE_TIME = timedelta(days=7)
|
|
|
|
def get_kroki_diagram(source, language):
|
|
encoded_diagram = base64.urlsafe_b64encode(zlib.compress(source.encode(), 9)).decode()
|
|
|
|
svg = download_and_cache_url(f"https://kroki.io/{language}/svg/{encoded_diagram}", DIAGRAM_CACHE_TIME)
|
|
|
|
return scourString(filter_svg(svg)).encode()
|
|
|
|
|
|
def fence_div_format(site_dir, source, language, *args, attrs, **kwargs):
|
|
try:
|
|
title = attrs.get("title", "")
|
|
|
|
diagram = get_kroki_diagram(source, language)
|
|
|
|
filename = hashlib.md5(source.encode()).hexdigest()[:12] + ".svg"
|
|
|
|
if title:
|
|
filename = "-".join([slugify(title)[:32].rstrip("-"), filename])
|
|
|
|
image_url = f"/_gen/kroki/{filename}"
|
|
|
|
write_file(diagram, os.path.join(site_dir, image_url.removeprefix("/")))
|
|
|
|
except Exception:
|
|
logger.exception("Failed to generate diagram")
|
|
return
|
|
|
|
return f"<img src='{image_url}' alt='{title}' />"
|
|
|
|
def on_config(config):
|
|
custom_fences = [
|
|
{
|
|
"name": diagram,
|
|
"class": diagram,
|
|
"format": partial(fence_div_format, config["site_dir"])
|
|
}
|
|
for diagram in DIAGRAM_TYPES
|
|
]
|
|
|
|
config.mdx_configs["pymdownx.superfences"] = {
|
|
"custom_fences": custom_fences
|
|
}
|
|
return config
|
|
|
|
|
|
def on_startup(command, dirty):
|
|
# Override user-agent so kroki accepts it
|
|
if _opener is None:
|
|
opener = build_opener()
|
|
install_opener(opener)
|
|
else:
|
|
opener = _opener
|
|
|
|
opener.addheaders = [('User-Agent','mkdocs')]
|