Improve code highlight

app/utils/highlight.py

@@ -3,6 +3,7 @@ from functools import lru_cache
 from bs4 import BeautifulSoup  # type: ignore
 from pygments import highlight as phighlight  # type: ignore
 from pygments.formatters import HtmlFormatter  # type: ignore
+from pygments.lexers import get_lexer_by_name  # type: ignore
 from pygments.lexers import guess_lexer  # type: ignore
 
 from app.config import CODE_HIGHLIGHTING_THEME
@@ -18,15 +19,29 @@ def highlight(html: str) -> str:
     for code in soup.find_all("code"):
         if not code.parent.name == "pre":
             continue
+
+        # Replace <br> tags with line breaks (Mastodon sends code like this)
         code_content = (
             code.encode_contents().decode().replace("<br>", "\n").replace("<br/>", "\n")
         )
-        lexer = guess_lexer(code_content)
-        tag = BeautifulSoup(
-            phighlight(code_content, lexer, _FORMATTER), "html5lib"
-        ).body.next
-        pre = code.parent
-        pre.replaceWith(tag)
-    out = soup.body
-    out.name = "div"
-    return str(out)
+
+        # 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-"):
+            try:
+                lexer = get_lexer_by_name(
+                    code.attrs["class"][0].removeprefix("language-")
+                )
+            except Exception:
+                lexer = guess_lexer(code_content)
+        else:
+            lexer = guess_lexer(code_content)
+
+        # Replace the code with Pygment output
+        code.parent.replaceWith(
+            BeautifulSoup(
+                phighlight(code_content, lexer, _FORMATTER), "html5lib"
+            ).body.next
+        )
+
+    return soup.body.encode_contents().decode()