refactor(rehype-momiji): use visit

taga3s · Sep 14, 2024 · b80a3fd · b80a3fd
1 parent 19b9176
commit b80a3fd
Showing 1 changed file with 31 additions and 21 deletions.
diff --git a/app/modules/rehype-momiji/buildCodeBlockHTML.ts b/app/modules/rehype-momiji/buildCodeBlockHTML.ts
@@ -3,6 +3,7 @@ import { getHighlighter } from "./highlighter";
 import { toHtml } from "hast-util-to-html";
 import type { Element } from "hast";
 import { COLOR_SHIRONERI, COLOR_HAI } from "./colors";
+import { visit } from "unist-util-visit";
 
 const defaultHighlighter = await getHighlighter({ themes: bundledThemes, langs: bundledLanguages });
 const defaultFilenameUIColor = `color: ${COLOR_SHIRONERI}; background-color: ${COLOR_HAI};`;
@@ -13,30 +14,39 @@ const buildCodeBlockHTML = (rawCode: string, lang: string, filename: string, the
     theme: theme,
   });
 
-  const targetElem = hast.children[0] as Element;
-
   // Add filename to the code block if it exists
-  if (filename !== "" && targetElem) {
-    targetElem.properties = {
-      class: targetElem.properties.class,
-      style: `${targetElem.properties.style}; padding-top: 0px;`,
-      tabindex: targetElem.properties.tabindex,
-    };
+  if (filename === "") {
+    return toHtml(hast);
+  }
+
+  visit(hast, "element", (node: Element) => {
+    if (
+      node.tagName === "pre" &&
+      node.properties.class &&
+      typeof node.properties.class === "string" &&
+      node.properties.class.includes("shiki")
+    ) {
+      node.properties = {
+        class: node.properties.class,
+        style: `${node.properties.style}; padding-top: 0px;`,
+        tabindex: node.properties.tabindex,
+      };
 
-    targetElem.children.unshift({
-      type: "element",
-      tagName: "div",
-      properties: {
-        style: `width: fit-content; margin-bottom: 16px; padding: 4px 8px; font-size: 14px; border-radius: 0 0 4px 4px; ${defaultFilenameUIColor};`,
-      },
-      children: [
-        {
-          type: "text",
-          value: filename,
+      node.children.unshift({
+        type: "element",
+        tagName: "div",
+        properties: {
+          style: `width: fit-content; margin-bottom: 16px; padding: 4px 8px; font-size: 14px; border-radius: 0 0 4px 4px; ${defaultFilenameUIColor};`,
         },
-      ],
-    });
-  }
+        children: [
+          {
+            type: "text",
+            value: filename,
+          },
+        ],
+      });
+    }
+  });
 
   return toHtml(hast);
 };