mkdocs-site/hooks/kroki.py

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')]