56 lines
1.6 KiB
Python
56 lines
1.6 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
|
||
|
|
import hashlib
|
||
|
|
from functools import partial
|
||
|
|
from mkdocs.utils import write_file
|
||
|
|
import os
|
||
|
|
|
||
|
|
logger = get_plugin_logger("kroki")
|
||
|
|
|
||
|
|
DIAGRAM_CACHE_TIME = timedelta(days=7)
|
||
|
|
|
||
|
|
def get_kroki_diagram(source, language):
|
||
|
|
encoded_diagram = base64.urlsafe_b64encode(zlib.compress(source.encode(), 9)).decode()
|
||
|
|
|
||
|
|
return download_and_cache_url(f"https://kroki.io/{language}/svg/{encoded_diagram}", DIAGRAM_CACHE_TIME)
|
||
|
|
|
||
|
|
|
||
|
|
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() + ".svg"
|
||
|
|
|
||
|
|
image_url = f"/assets/kroki-generated/{language}/{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):
|
||
|
|
# Override user-agent so kroki accepts it
|
||
|
|
opener = build_opener()
|
||
|
|
opener.addheaders = [('User-Agent','mkdocs')]
|
||
|
|
install_opener(opener)
|
||
|
|
|
||
|
|
custom_fences = [{
|
||
|
|
"name": "blockdiag",
|
||
|
|
"class": "blockdiag",
|
||
|
|
"format": partial(fence_div_format, config["site_dir"])
|
||
|
|
}]
|
||
|
|
|
||
|
|
config.mdx_configs["pymdownx.superfences"] = {
|
||
|
|
"custom_fences": custom_fences
|
||
|
|
}
|
||
|
|
return config
|