feat: support multiple transformers

Author: antfu

README.md

@@ -435,19 +435,21 @@ Since `shikiji` uses `hast` internally, you can use the `transforms` option to c
 const code = await codeToHtml('foo\bar', {
   lang: 'js',
   theme: 'vitesse-light',
-  transforms: {
-    code(node) {
-      node.properties.class = 'language-js'
-    },
-    line(node, line) {
-      node.properties['data-line'] = line
-      if ([1, 3, 4].includes(line))
-        node.properties.class += ' highlight'
-    },
-    token(node, line, col) {
-      node.properties.class = `token:${line}:${col}`
+  transformers: [
+    {
+      code(node) {
+        node.properties.class = 'language-js'
+      },
+      line(node, line) {
+        node.properties['data-line'] = line
+        if ([1, 3, 4].includes(line))
+          node.properties.class += ' highlight'
+      },
+      token(node, line, col) {
+        node.properties.class = `token:${line}:${col}`
+      },
     },
-  },
+  ]
 })
 ```
 

packages/markdown-it-shikiji/src/index.ts

@@ -32,7 +32,7 @@ function setup(markdownit: MarkdownIt, highlighter: Highlighter, options: Markdo
       lang,
     }
 
-    codeOptions.transforms ||= {}
+    codeOptions.transformers ||= []
 
     if (highlightLines) {
       const lines = parseHighlightLines(attrs)
@@ -41,17 +41,23 @@ function setup(markdownit: MarkdownIt, highlighter: Highlighter, options: Markdo
           ? 'highlighted'
           : highlightLines
 
-        codeOptions.transforms.line = (node, line) => {
-          if (lines.includes(line))
-            node.properties.class += ` ${className}`
-          return node
-        }
+        codeOptions.transformers.push({
+          name: 'markdown-it-shikiji:line-class',
+          line(node, line) {
+            if (lines.includes(line))
+              node.properties.class += ` ${className}`
+            return node
+          },
+        })
       }
     }
 
-    codeOptions.transforms.code = (node) => {
-      node.properties.class = `language-${lang}`
-    }
+    codeOptions.transformers.push({
+      name: 'markdown-it-shikiji:block-class',
+      code(node) {
+        node.properties.class = `language-${lang}`
+      },
+    })
 
     return highlighter.codeToHtml(
       code,

packages/rehype-shikiji/src/index.ts

@@ -100,17 +100,18 @@ const rehypeShikiji: Plugin<[RehypeShikijiOptions], Root> = function (options =
             ? 'highlighted'
             : highlightLines
 
-          codeOptions.transforms ||= {}
-          codeOptions.transforms.line = (node, line) => {
-            if (lines.includes(line))
-              node.properties.class += ` ${className}`
-            return node
-          }
+          codeOptions.transformers ||= []
+          codeOptions.transformers.push({
+            name: 'rehype-shikiji:line-class',
+            line(node, line) {
+              if (lines.includes(line))
+                node.properties.class += ` ${className}`
+              return node
+            },
+          })
         }
       }
-
       const fragment = highlighter.codeToHast(code, codeOptions)
-
       parent.children.splice(index, 1, ...fragment.children)
     })
   }

packages/shikiji-compat/src/index.ts

@@ -65,17 +65,18 @@ export async function getHighlighter(options: HighlighterOptions = {}) {
     }
 
     if (options.lineOptions) {
-      options.transforms = options.transforms || {}
-      const prev = options.transforms.line
-      options.transforms.line = (ast, line) => {
-        const node = prev?.(ast, line) || ast
-        const lineOption = options.lineOptions?.find(o => o.line === line)
-        if (lineOption?.classes) {
-          node.properties ??= {}
-          node.properties.class = [node.properties.class, ...lineOption.classes].filter(Boolean).join(' ')
-        }
-        return node
-      }
+      options.transformers ||= []
+      options.transformers.push({
+        name: 'shikiji-compat:line-class',
+        line(node, line) {
+          const lineOption = options.lineOptions?.find(o => o.line === line)
+          if (lineOption?.classes) {
+            node.properties ??= {}
+            node.properties.class = [node.properties.class, ...lineOption.classes].filter(Boolean).join(' ')
+          }
+          return node
+        },
+      })
     }
 
     return shikiji.codeToHtml(code, options as any)

packages/shikiji/src/core/renderer-hast.ts

@@ -114,8 +114,15 @@ export function tokensToHast(
 ) {
   const {
     mergeWhitespaces = true,
+    transformers = [],
   } = options
 
+  // TODO: remove this in next major version
+  if (options.transforms) {
+    transformers.push(options.transforms)
+    console.warn('[shikiji] `transforms` option is deprecated, use `transformers` instead')
+  }
+
   if (mergeWhitespaces)
     tokens = mergeWhitespaceTokens(tokens)
 
@@ -169,23 +176,33 @@ export function tokensToHast(
       if (style)
         tokenNode.properties.style = style
 
-      tokenNode = options.transforms?.token?.(tokenNode, idx + 1, col, lineNode) || tokenNode
+      for (const transformer of transformers)
+        tokenNode = transformer?.token?.(tokenNode, idx + 1, col, lineNode) || tokenNode
 
       lineNode.children.push(tokenNode)
       col += token.content.length
     }
 
-    lineNode = options.transforms?.line?.(lineNode, idx + 1) || lineNode
+    for (const transformer of transformers)
+      lineNode = transformer?.line?.(lineNode, idx + 1) || lineNode
+
     lines.push(lineNode)
   })
 
-  codeNode = options.transforms?.code?.(codeNode) || codeNode
+  for (const transformer of transformers)
+    codeNode = transformer?.code?.(codeNode) || codeNode
   preNode.children.push(codeNode)
 
-  preNode = options.transforms?.pre?.(preNode) || preNode
+  for (const transformer of transformers)
+    preNode = transformer?.pre?.(preNode) || preNode
+
   tree.children.push(preNode)
 
-  return options.transforms?.root?.(tree) || tree
+  let result = tree
+  for (const transformer of transformers)
+    result = transformer?.root?.(result) || result
+
+  return result
 }
 
 function getTokenStyles(token: ThemedToken) {

packages/shikiji/src/types.ts

@@ -128,7 +128,12 @@ export interface CodeToHastOptionsCommon<Languages extends string = string> {
   /**
    * Transform the generated HAST tree.
    */
-  transforms?: HastTransformers
+  transformers?: ShikijiTransformer[]
+
+  /**
+   * @deprecated use `transformers` instead
+   */
+  transforms?: ShikijiTransformer
 }
 
 export interface CodeToTokensWithThemesOptions<Languages = string, Themes = string> {
@@ -266,7 +271,11 @@ export interface ThemeRegistration extends ThemeRegistrationRaw {
   colors?: Record<string, string>
 }
 
-export interface HastTransformers {
+export interface ShikijiTransformer {
+  /**
+   * Name of the transformer
+   */
+  name?: string
   /**
    * Transform the entire generated HAST tree. Return a new Node will replace the original one.
    *
@@ -291,10 +300,12 @@ export interface HtmlRendererOptionsCommon {
   fg?: string
   bg?: string
 
+  transformers?: ShikijiTransformer[]
+
   /**
-   * Hast transformers
+   * @deprecated use `transformers` instead
    */
-  transforms?: HastTransformers
+  transforms?: ShikijiTransformer
 
   themeName?: string
 

packages/shikiji/test/hast.test.ts

@@ -1,7 +1,12 @@
-import { describe, expect, it } from 'vitest'
+/* eslint-disable style/no-tabs */
+import { afterEach, describe, expect, it, vi } from 'vitest'
 import { toHtml } from 'hast-util-to-html'
 import { codeToHtml, getHighlighter } from '../src'
 
+afterEach(() => {
+  vi.restoreAllMocks
+})
+
 describe('should', () => {
   it('works', async () => {
     const shiki = await getHighlighter({
@@ -19,9 +24,13 @@ describe('should', () => {
   })
 
   it('hast transformer', async () => {
+    const warn = vi.spyOn(console, 'warn')
+    warn.mockImplementation((() => {}) as any)
+
     const code = await codeToHtml('foo\bar', {
       lang: 'js',
       theme: 'vitesse-light',
+      // Use deprecated `transforms` option on purpose to test
       transforms: {
         line(node, line) {
           node.properties['data-line'] = line
@@ -37,6 +46,10 @@ describe('should', () => {
 
     expect(code)
       .toMatchInlineSnapshot('"<pre class=\\"shiki vitesse-light\\" style=\\"background-color:#ffffff;color:#393a34\\" tabindex=\\"0\\"><code class=\\"language-js\\"><span class=\\"line\\" data-line=\\"1\\"><span style=\\"color:#B07D48\\" class=\\"token:1:0\\">foo</span><span style=\\"color:#393A34\\" class=\\"token:1:3\\"></span><span style=\\"color:#B07D48\\" class=\\"token:1:4\\">ar</span></span></code></pre>"')
+
+    expect(warn).toBeCalledTimes(1)
+    expect(warn.mock.calls[0][0])
+      .toMatchInlineSnapshot('"[shikiji] `transforms` option is deprecated, use `transformers` instead"')
   })
 })
 
@@ -48,19 +61,21 @@ it('hasfocus support', async () => {
   const code = await codeToHtml(snippet, {
     lang: 'php',
     theme: 'vitesse-light',
-    transforms: {
-      code(node) {
-        node.properties.class = 'language-php'
-      },
-      token(node, line, col, parent) {
-        node.children.forEach((child) => {
-          if (child.type === 'text' && child.value.includes('[!code focus]')) {
-            parent.properties['data-has-focus'] = 'true'
-            node.children.splice(node.children.indexOf(child), 1)
-          }
-        })
+    transformers: [
+      {
+        code(node) {
+          node.properties.class = 'language-php'
+        },
+        token(node, line, col, parent) {
+          node.children.forEach((child) => {
+            if (child.type === 'text' && child.value.includes('[!code focus]')) {
+              parent.properties['data-has-focus'] = 'true'
+              node.children.splice(node.children.indexOf(child), 1)
+            }
+          })
+        },
       },
-    },
+    ],
   })
 
   expect(code)
@@ -70,3 +85,59 @@ it('hasfocus support', async () => {
       <span class=\\"line\\"><span style=\\"color:#999999\\">$</span><span style=\\"color:#B07D48\\">bar</span><span style=\\"color:#999999\\"> =</span><span style=\\"color:#B5695999\\"> \\"</span><span style=\\"color:#B56959\\">baz</span><span style=\\"color:#B5695999\\">\\"</span><span style=\\"color:#999999\\">;</span></span></code></pre>"
     `)
 })
+
+it('render whitespace', async () => {
+  const snippet = [
+    '  space()',
+    '\t\ttab()',
+  ].join('\n')
+
+  const classMap: Record<string, string> = {
+    ' ': 'space',
+    '\t': 'tab',
+  }
+
+  const code = await codeToHtml(snippet, {
+    lang: 'js',
+    theme: 'vitesse-light',
+    transformers: [
+      {
+        line(node) {
+          const first = node.children[0]
+          if (!first || first.type !== 'element')
+            return
+          const textNode = first.children[0]
+          if (!textNode || textNode.type !== 'text')
+            return
+          node.children = node.children.flatMap((child) => {
+            if (child.type !== 'element')
+              return child
+            const node = child.children[0]
+            if (node.type !== 'text' || !node.value)
+              return child
+            const parts = node.value.split(/([ \t])/).filter(i => i.length)
+            if (parts.length <= 1)
+              return child
+
+            return parts.map((part) => {
+              const clone = {
+                ...child,
+                properties: { ...child.properties },
+              }
+              clone.children = [{ type: 'text', value: part }]
+              if (part in classMap)
+                clone.properties.class = [clone.properties.class, classMap[part]].filter(Boolean).join(' ')
+              return clone
+            })
+          })
+        },
+      },
+    ],
+  })
+
+  expect(code)
+    .toMatchInlineSnapshot(`
+      "<pre class=\\"shiki vitesse-light\\" style=\\"background-color:#ffffff;color:#393a34\\" tabindex=\\"0\\"><code><span class=\\"line\\"><span style=\\"color:#59873A\\" class=\\"space\\"> </span><span style=\\"color:#59873A\\" class=\\"space\\"> </span><span style=\\"color:#59873A\\">space</span><span style=\\"color:#999999\\">()</span></span>
+      <span class=\\"line\\"><span style=\\"color:#59873A\\" class=\\"tab\\">	</span><span style=\\"color:#59873A\\" class=\\"tab\\">	</span><span style=\\"color:#59873A\\">tab</span><span style=\\"color:#999999\\">()</span></span></code></pre>"
+    `)
+})