Build my own kroki integration

It's faster, easier to work with, and properly supports superfences

content/other.md

@@ -1,6 +1,6 @@
 # Other
 
-```kroki-blockdiag no-transparency=false
+```{.blockdiag title="Blocks"}
 blockdiag {
   blockdiag -> generates -> "block-diagrams";
   blockdiag -> is -> "very easy!";

hooks/kroki.py

@@ -0,0 +1,55 @@
+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

mkdocs.yml

@@ -17,6 +17,7 @@ hooks:
   - hooks/redirects.py
   - hooks/globals.py
   - hooks/pygments-styles.py
+  - hooks/kroki.py
 
 markdown_extensions:
   - attr_list
@@ -24,7 +25,4 @@ markdown_extensions:
   - pymdownx.superfences
 
 plugins:
-  # - kroki:
-  #     DownloadImages: true
-  #     DownloadDir: assets/kroki-generated
   - search

requirements.txt

@@ -1,6 +1,5 @@
 mkdocs
 jinja2-simple-tags
-mkdocs-kroki-plugin
 pymdown-extensions
 pygments
 minify-html