From 7b784e3011d2e79393b6c1f8e8070c2a59e44775 Mon Sep 17 00:00:00 2001
From: Thomas Sileo <t@a4.io>
Date: Fri, 6 Jan 2023 21:21:53 +0100
Subject: [PATCH] Tweak code highlight

---
 app/source.py          | 16 +++++++++++-----
 app/utils/highlight.py | 27 +++++++++++++--------------
 2 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/app/source.py b/app/source.py
index a46707f..20e98ee 100644
--- a/app/source.py
+++ b/app/source.py
@@ -3,12 +3,12 @@ import typing
 
 from loguru import logger
 from mistletoe import Document  # type: ignore
+from mistletoe.block_token import CodeFence  # type: ignore
 from mistletoe.html_renderer import HTMLRenderer  # type: ignore
 from mistletoe.span_token import SpanToken  # type: ignore
-from pygments import highlight  # type: ignore
 from pygments.formatters import HtmlFormatter  # type: ignore
 from pygments.lexers import get_lexer_by_name as get_lexer  # type: ignore
-from pygments.lexers import guess_lexer  # type: ignore
+from pygments.util import ClassNotFound  # type: ignore
 from sqlalchemy import select
 
 from app import webfinger
@@ -104,10 +104,16 @@ class CustomRenderer(HTMLRenderer):
         )
         return link
 
-    def render_block_code(self, token: typing.Any) -> str:
+    def render_block_code(self, token: CodeFence) -> str:
+        lexer_attr = ""
+        try:
+            lexer = get_lexer(token.language)
+            lexer_attr = f' data-microblogpub-lexer="{lexer.aliases[0]}"'
+        except ClassNotFound:
+            pass
+
         code = token.children[0].content
-        lexer = get_lexer(token.language) if token.language else guess_lexer(code)
-        return highlight(code, lexer, _FORMATTER)
+        return f"<pre><code{lexer_attr}>\n{code}\n</code></pre>"
 
 
 async def _prefetch_mentioned_actors(
diff --git a/app/utils/highlight.py b/app/utils/highlight.py
index ff0e079..2e60a36 100644
--- a/app/utils/highlight.py
+++ b/app/utils/highlight.py
@@ -32,23 +32,22 @@ def highlight(html: str) -> str:
 
         # If this comes from a microblog.pub instance we may have the language
         # in the class name
-        if "class" in code.attrs and code.attrs["class"][0].startswith("language-"):
+        if "data-microblogpub-lexer" in code.attrs:
             try:
-                lexer = get_lexer_by_name(
-                    code.attrs["class"][0].removeprefix("language-")
-                )
+                lexer = get_lexer_by_name(code.attrs["data-microblogpub-lexer"])
             except Exception:
                 lexer = guess_lexer(code_content)
-        else:
-            lexer = guess_lexer(code_content)
 
-        # Replace the code with Pygment output
-        # XXX: the HTML escaping causes issue with Python type annotations
-        code_content = code_content.replace(") -&gt; ", ") -> ")
-        code.parent.replaceWith(
-            BeautifulSoup(
-                phighlight(code_content, lexer, _FORMATTER), "html5lib"
-            ).body.next
-        )
+            # Replace the code with Pygment output
+            # XXX: the HTML escaping causes issue with Python type annotations
+            code_content = code_content.replace(") -&gt; ", ") -> ")
+            code.parent.replaceWith(
+                BeautifulSoup(
+                    phighlight(code_content, lexer, _FORMATTER), "html5lib"
+                ).body.next
+            )
+        else:
+            code.name = "div"
+            code["class"] = code.get("class", []) + ["highlight"]
 
     return soup.body.encode_contents().decode()