add code highlighting

app/lib/babycode.py

@@ -1,5 +1,9 @@
 from .babycode_parser import Parser
 from markupsafe import Markup, escape
+from pygments import highlight
+from pygments.formatters import HtmlFormatter
+from pygments.lexers import get_lexer_by_name
+from pygments.util import ClassNotFound as PygmentsClassNotFound
 import re
 
 BABYCODE_VERSION = 4
@@ -56,9 +60,17 @@ def tag_code(children, attr, surrounding):
     if is_inline:
         return f"<code class=\"inline-code\">{children}</code>"
     else:
-        t = children.strip()
-        button = f"<button type=button class=\"copy-code\" value=\"{t}\">Copy</button>"
-        return f"<pre><span class=\"copy-code-container\">{button}</span><code>{t}</code></pre>"
+        input_code = children.strip()
+        button = f"<button type=button class=\"copy-code\" value=\"{input_code}\">Copy</button>"
+        unhighlighted = f"<pre><span class=\"copy-code-container\"><span class=\"code-language-identifier\">code block</span>{button}</span><code>{input_code}</code></pre>"
+        if not attr:
+            return unhighlighted
+        try:
+            lexer = get_lexer_by_name(attr.strip())
+            formatter = HtmlFormatter(nowrap=True)
+            return f"<pre><span class=\"copy-code-container\"><span class=\"code-language-identifier\">{lexer.name}</span>{button}</span><code>{highlight(input_code.unescape(), lexer, formatter)}</code></pre>"
+        except PygmentsClassNotFound:
+            return unhighlighted
 
 def tag_list(children):
     list_body = re.sub(r"  +\n", "<br>", children.strip())