feat(markdown): support multiple themes for code highlighter (#1251)

Co-authored-by: Yaël Guilloux <yael.guilloux@gmail.com>

docs/content/4.api/3.configuration.md

@@ -202,9 +202,37 @@ Nuxt Content uses [Shiki](https://github.com/shikijs/shiki) to provide syntax hi
 
 | Option | Type | Description |
 | ----------------- | :--------: | :-------- |
-| `theme`     | `ShikiTheme` | The [color theme](https://github.com/shikijs/shiki/blob/main/docs/themes.md) to use |
+| `theme`     | `ShikiTheme` or `Record<string, ShikiTheme>` | The [color theme](https://github.com/shikijs/shiki/blob/main/docs/themes.md) to use. |
 | `preload`     | `ShikiLang[]` | The [preloaded languages](https://github.com/shikijs/shiki/blob/main/docs/languages.md) available for highlighting. |
 
+#### `highlight.theme`
+
+Theme can be specified by a single string but also supports an object with multiple themes.
+
+This option is compatible with [Color Mode module](https://color-mode.nuxtjs.org/).
+
+If you are using multiple themes, it's recommended to always have a `default` theme specified.
+
+```ts
+export default defineNuxtConfig({
+  content: {
+    highlight: {
+      // Theme used in all color schemes.
+      theme: 'github-light'
+      // OR
+      theme: {
+        // Default theme (same as single string)
+        default: 'github-light',
+        // Theme used if `html.dark`
+        dark: 'github-dark'
+        // Theme used if `html.sepia`
+        sepia: 'monokai'
+      }
+    }
+  }
+})
+```
+
 ## `yaml`
 
 - Type: `false | Object`{lang=ts}

src/module.ts

@@ -122,7 +122,10 @@ export interface ModuleOptions {
     /**
      * Default theme that will be used for highlighting code blocks.
      */
-    theme?: ShikiTheme,
+    theme?: ShikiTheme | {
+      default: ShikiTheme
+      [theme: string]: ShikiTheme
+    },
     /**
      * Preloaded languages that will be available for highlighting code blocks.
      */

src/runtime/server/api/highlight.ts

@@ -15,13 +15,26 @@ const resolveLang = (lang: string): Lang | undefined =>
 /**
  * Resolve Shiki compatible theme from string.
  */
-const resolveTheme = (theme: string): Theme | undefined =>
-  BUNDLED_THEMES.find(t => t === theme)
+const resolveTheme = (theme: string | Record<string, string>): Record<string, Theme> | undefined => {
+  if (!theme) {
+    return
+  }
+  if (typeof theme === 'string') {
+    theme = {
+      default: theme
+    }
+  }
+
+  return Object.entries(theme).reduce((acc, [key, value]) => {
+    acc[key] = BUNDLED_THEMES.find(t => t === value)
+    return acc
+  }, {})
+}
 
 /**
  * Resolve Shiki highlighter compatible payload from request body.
  */
-const resolveBody = (body: Partial<HighlightParams>): { code: string, lang?: Lang, theme?: Theme } => {
+const resolveBody = (body: Partial<HighlightParams>) => {
   // Assert body schema
   if (typeof body.code !== 'string') { throw createError({ statusMessage: 'Bad Request', statusCode: 400, message: 'Missing code key.' }) }
 
@@ -40,7 +53,7 @@ export default defineLazyEventHandler(async () => {
 
   // Initialize highlighter with defaults
   const highlighter = await getHighlighter({
-    theme: theme || 'dark-plus',
+    theme: theme?.default || theme || 'dark-plus',
     langs: [
       ...(preload || ['json', 'js', 'ts', 'css']),
       'shell',
@@ -60,7 +73,7 @@ export default defineLazyEventHandler(async () => {
   return async (event): Promise<HighlightThemedToken[][]> => {
     const params = await useBody<Partial<HighlightParams>>(event)
 
-    const { code, lang, theme } = resolveBody(params)
+    const { code, lang, theme = { default: highlighter.getTheme() } } = resolveBody(params)
 
     // Skip highlight if lang is not supported
     if (!lang) {
@@ -73,21 +86,88 @@ export default defineLazyEventHandler(async () => {
     }
 
     // Load supported theme on-demand
-    if (theme && !highlighter.getLoadedThemes().includes(theme)) {
-      await highlighter.loadTheme(theme)
-    }
+    await Promise.all(
+      Object.values(theme).map(async (theme) => {
+        if (!highlighter.getLoadedThemes().includes(theme)) {
+          await highlighter.loadTheme(theme)
+        }
+      })
+    )
 
     // Highlight code
-    const highlightedCode = highlighter.codeToThemedTokens(code, lang, theme)
-
-    // Clean up to shorten response payload
-    for (const line of highlightedCode) {
-      for (const token of line) {
-        delete token.fontStyle
-        delete token.explanation
+    const coloredTokens = Object.entries(theme).map(([key, theme]) => {
+      const tokens = highlighter.codeToThemedTokens(code, lang, theme)
+      return {
+        key,
+        theme,
+        tokens
       }
+    })
+
+    const highlightedCode: HighlightThemedToken[][] = []
+    for (const line in coloredTokens[0].tokens) {
+      highlightedCode[line] = coloredTokens.reduce((acc, color) => {
+        return mergeLines({
+          key: coloredTokens[0].key,
+          tokens: acc
+        }, {
+          key: color.key,
+          tokens: color.tokens[line]
+        })
+      }, coloredTokens[0].tokens[line])
     }
 
     return highlightedCode
   }
 })
+
+function mergeLines (line1, line2) {
+  const mergedTokens = []
+  const getColors = (h, i) => typeof h.tokens[i].color === 'string' ? { [h.key]: h.tokens[i].color } : h.tokens[i].color
+
+  const [big, small] = line1.tokens.length > line2.tokens.length ? [line1, line2] : [line2, line1]
+  let targetToken = 0
+  let targetTokenCharIndex = 0
+  big.tokens.forEach((t, i) => {
+    if (targetTokenCharIndex === 0) {
+      if (t.content === small.tokens[i]?.content) {
+        mergedTokens.push({
+          content: t.content,
+          color: {
+            ...getColors(big, i),
+            ...getColors(small, i)
+          }
+        })
+        targetToken = i + 1
+        return
+      }
+      if (t.content === small.tokens[targetToken]?.content) {
+        mergedTokens.push({
+          content: t.content,
+          color: {
+            ...getColors(big, i),
+            ...getColors(small, targetToken)
+          }
+        })
+        targetToken += 1
+        return
+      }
+    }
+
+    if (small.tokens[targetToken]?.content?.substring(targetTokenCharIndex, targetTokenCharIndex + t.content.length) === t.content) {
+      targetTokenCharIndex += t.content.length
+      mergedTokens.push({
+        content: t.content,
+        color: {
+          ...getColors(big, i),
+          ...getColors(small, targetToken)
+        }
+      })
+    }
+    if (small.tokens[targetToken]?.content.length <= targetTokenCharIndex) {
+      targetToken += 1
+      targetTokenCharIndex = 0
+    }
+  })
+  return mergedTokens
+}

src/runtime/server/transformers/shiki.ts

@@ -2,6 +2,8 @@ import { visit } from 'unist-util-visit'
 import { withBase } from 'ufo'
 import { useRuntimeConfig } from '#imports'
 
+const highlightConfig = useRuntimeConfig().content.highlight
+
 const withContentBase = (url: string) => {
   return withBase(url, `/api/${useRuntimeConfig().public.content.base}`)
 }
@@ -10,74 +12,122 @@ export default {
   name: 'markdown',
   extensions: ['.md'],
   transform: async (content) => {
-    const codeBlocks = []
+    const tokenColors: Record<string, {colors: any, className: string}> = {}
+    const codeBlocks: any[] = []
+    const inlineCodes: any = []
     visit(
       content.body,
-      (node: any) => node.tag === 'code' && node?.props.code,
-      (node) => { codeBlocks.push(node) }
+      (node: any) => (node.tag === 'code' && node?.props.code) || (node.tag === 'code-inline' && (node.props?.lang || node.props?.language)),
+      (node) => {
+        if (node.tag === 'code') {
+          codeBlocks.push(node)
+        } else if (node.tag === 'code-inline') {
+          inlineCodes.push(node)
+        }
+      }
     )
+
     await Promise.all(codeBlocks.map(highlightBlock))
+    await Promise.all(inlineCodes.map(highlightInline))
 
-    const inlineCodes = []
-    visit(
-      content.body,
-      (node: any) => node.tag === 'code-inline' && (node.props?.lang || node.props?.language),
-      (node) => { inlineCodes.push(node) }
-    )
+    // Inject token colors at the end of the document
+    if (Object.values(tokenColors).length) {
+      const colors: string[] = []
+      for (const colorClass of Object.values(tokenColors)) {
+        Object.entries(colorClass.colors).forEach(([variant, color]) => {
+          if (variant === 'default') {
+            colors.unshift(`.${colorClass.className}{color:${color}}`)
+          } else {
+            colors.push(`.${variant} .${colorClass.className}{color:${color}}`)
+          }
+        })
+      }
 
-    await Promise.all(inlineCodes.map(highlightInline))
+      content.body.children.push({
+        type: 'element',
+        tag: 'style',
+        children: [{ type: 'text', value: colors.join('') }]
+      })
+    }
 
     return content
-  }
-}
 
-const tokenSpan = ({ content, color }) => ({
-  type: 'element',
-  tag: 'span',
-  props: { style: { color } },
-  children: [{ type: 'text', value: content }]
-})
-
-const highlightInline = async (node) => {
-  const code = node.children[0].value
-
-  // Fetch highlighted tokens
-  const lines = await $fetch(withContentBase('highlight'), {
-    method: 'POST',
-    body: {
-      code,
-      lang: node.props.lang || node.props.language
+    /**
+     * Highlight inline code
+     */
+    async function highlightInline (node) {
+      const code = node.children[0].value
+
+      // Fetch highlighted tokens
+      const lines = await $fetch<any[]>(withContentBase('highlight'), {
+        method: 'POST',
+        body: {
+          code,
+          lang: node.props.lang || node.props.language,
+          theme: highlightConfig.theme
+        }
+      })
+
+      // Generate highlighted children
+      node.children = lines[0].map(tokenSpan)
+
+      node.props = node.props || {}
+      node.props.class = 'colored'
+
+      return node
     }
-  })
 
-  // Generate highlighted children
-  node.children = lines[0].map(tokenSpan)
+    /**
+     * Highlight a code block
+     */
+    async function highlightBlock (node) {
+      const { code, language: lang, highlights = [] } = node.props
 
-  node.props = node.props || {}
-  node.props.class = 'colored'
+      // Fetch highlighted tokens
+      const lines = await $fetch<any[]>(withContentBase('highlight'), {
+        method: 'POST',
+        body: {
+          code,
+          lang,
+          theme: highlightConfig.theme
+        }
+      })
 
-  return node
-}
+      // Generate highlighted children
+      const innerCodeNode = node.children[0].children[0]
+      innerCodeNode.children = lines.map((line, lineIndex) => ({
+        type: 'element',
+        tag: 'span',
+        props: { class: ['line', highlights.includes(lineIndex + 1) ? 'highlight' : ''].join(' ').trim() },
+        children: line.map(tokenSpan)
+      }))
+      return node
+    }
 
-const highlightBlock = async (node) => {
-  const { code, language: lang, highlights = [] } = node.props
+    function getColorProps (token) {
+      if (!token.color) {
+        return {}
+      }
+      if (typeof token.color === 'string') {
+        return { style: { color: token.color } }
+      }
+      const key = Object.values(token.color).join('')
+      if (!tokenColors[key]) {
+        tokenColors[key] = {
+          colors: token.color,
+          className: 'ct-' + Math.random().toString(16).substring(2, 8) // hash(key)
+        }
+      }
+      return { class: tokenColors[key].className }
+    }
 
-  // Fetch highlighted tokens
-  const lines = await $fetch(withContentBase('highlight'), {
-    method: 'POST',
-    body: {
-      code,
-      lang
+    function tokenSpan (token) {
+      return {
+        type: 'element',
+        tag: 'span',
+        props: getColorProps(token),
+        children: [{ type: 'text', value: token.content }]
+      }
     }
-  })
-
-  // Generate highlighted children
-  const innerCodeNode = node.children[0].children[0]
-  innerCodeNode.children = lines.map((line, lineIndex) => ({
-    type: 'element',
-    tag: 'span',
-    props: { class: ['line', highlights.includes(lineIndex + 1) ? 'highlight' : ''].join(' ').trim() },
-    children: line.map(tokenSpan)
-  }))
-  return node
+  }
 }

src/runtime/types.d.ts

@@ -283,10 +283,10 @@ export interface NavItem {
 export interface HighlightParams {
   code: string
   lang: string
-  theme: Theme
+  theme: Theme | Record<string, Theme>
 }
 
 export interface HighlightThemedToken {
   content: string
-  color?: string
+  color?: string | Record<string, string>
 }