feat!: rework hastTransform to transforms

antfu · antfu · commit 3d2459126010 · 2023-08-21T15:00:00.000+02:00
diff --git a/README.md b/README.md
@@ -415,6 +415,30 @@ console.log(root)
 }
 ```
 
+### Hast transformers
+
+Since `shikiji` uses `hast` internally, you can use the `transforms` option to customize the generated HTML by manipulating the hast tree. You can pass custom functions to modify the tree for different types of nodes. For example:
+
+```js
+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}`
+    },
+  },
+})
+```
+
 ## Breaking Changes from Shiki
 
 As of [`shiki@0.4.3`](https://github.com/shikijs/shiki/releases/tag/v0.14.3):
@@ -425,7 +449,7 @@ As of [`shiki@0.4.3`](https://github.com/shikijs/shiki/releases/tag/v0.14.3):
 - Highlighter does not maintain an internal default theme context. `theme` option is required for `codeToHtml` and `codeToThemedTokens`.
 - `.ansiToHtml` is merged into `.codeToHtml` as a special language `ansi`. Use `.codeToHtml(code, { lang: 'ansi' })` instead.
 - `codeToHtml` uses [`hast`](https://github.com/syntax-tree/hast) internally. The generated HTML will be a bit different but should behavior the same.
-- `lineOptions` is dropped in favor of the fully customizable `hastTransform` option.
+- `lineOptions` is dropped in favor of the fully customizable `transforms` option.
 - CJS and IIFE builds are dropped. See [CJS Usage](#cjs-usage) and [CDN Usage](#cdn-usage) for more details.
 - `LanguageRegistration`'s `grammar` field is flattened to `LanguageRegistration` itself (refer to the types for more details).
 
diff --git a/packages/shikiji/src/core/renderer-hast.ts b/packages/shikiji/src/core/renderer-hast.ts
@@ -99,38 +99,39 @@ export function tokensToHast(
   const lines: (Element | Text)[] = []
   const tree: Root = {
     type: 'root',
-    children: [
-      {
-        type: 'element',
-        tagName: 'pre',
-        properties: {
-          class: `shiki ${options.themeName || ''}`,
-          style: options.rootStyle || `background-color:${options.bg};color:${options.fg}`,
-          tabindex: '0',
-        },
-        children: [
-          {
-            type: 'element',
-            tagName: 'code',
-            properties: {},
-            children: lines,
-          },
-        ],
-      },
-    ],
+    children: [],
+  }
+
+  let preNode: Element = {
+    type: 'element',
+    tagName: 'pre',
+    properties: {
+      class: `shiki ${options.themeName || ''}`,
+      style: options.rootStyle || `background-color:${options.bg};color:${options.fg}`,
+      tabindex: '0',
+    },
+    children: [],
+  }
+
+  let codeNode: Element = {
+    type: 'element',
+    tagName: 'code',
+    properties: {},
+    children: lines,
   }
 
   tokens.forEach((line, idx) => {
     if (idx)
       lines.push({ type: 'text', value: '\n' })
 
-    const lineNode: Element = {
+    let lineNode: Element = {
       type: 'element',
       tagName: 'span',
       properties: { class: 'line' },
       children: [],
     }
-    lines.push(lineNode)
+
+    let col = 0
 
     for (const token of line) {
       const styles = [token.htmlStyle || `color:${token.color}`]
@@ -143,7 +144,7 @@ export function tokensToHast(
           styles.push('text-decoration:underline')
       }
 
-      const tokenNode: Element = {
+      let tokenNode: Element = {
         type: 'element',
         tagName: 'span',
         properties: {
@@ -152,11 +153,23 @@ export function tokensToHast(
         children: [{ type: 'text', value: token.content }],
       }
 
+      tokenNode = options.transforms?.token?.(tokenNode, idx + 1, col) || tokenNode
+
       lineNode.children.push(tokenNode)
+      col += token.content.length
     }
+
+    lineNode = options.transforms?.line?.(lineNode, idx + 1) || lineNode
+    lines.push(lineNode)
   })
 
-  return options.hastTransform?.(tree) || tree
+  codeNode = options.transforms?.code?.(codeNode) || codeNode
+  preNode.children.push(codeNode)
+
+  preNode = options.transforms?.pre?.(preNode) || preNode
+  tree.children.push(preNode)
+
+  return options.transforms?.root?.(tree) || tree
 }
 
 function mergeWhitespaceTokens(tokens: ThemedToken[][]) {
diff --git a/packages/shikiji/src/types.ts b/packages/shikiji/src/types.ts
@@ -1,5 +1,5 @@
 import type { IGrammar, IRawGrammar, IRawTheme } from 'vscode-textmate'
-import type { Root } from 'hast'
+import type { Element, Root } from 'hast'
 import type { bundledThemes } from './themes'
 import type { bundledLanguages } from './assets/langs'
 import type { FontStyle } from './core/stackElementMetadata'
@@ -121,9 +121,9 @@ export interface CodeToThemedTokensOptions<Languages = string, Themes = string>
 export interface CodeToHastOptionsCommon<Languages = string> {
   lang: Languages | SpecialLanguage
   /**
-   * TODO
+   * Transform the generated HAST tree.
    */
-  hastTransform?: (hast: Root) => Root | void
+  transforms?: HastTransformers
 }
 
 export interface CodeToTokensWithThemesOptions<Languages = string, Themes = string> {
@@ -254,14 +254,34 @@ export interface ThemeRegistration extends ThemeRegistrationRaw {
   colors?: Record<string, string>
 }
 
+export interface HastTransformers {
+  /**
+   * Transform the entire generated HAST tree. Return a new Node will replace the original one.
+   *
+   * @param hast
+   */
+  root?: (hast: Root) => Root | void
+  pre?: (hast: Element) => Element | void
+  code?: (hast: Element) => Element | void
+  /**
+   * Transform each line element.
+   *
+   * @param hast
+   * @param line 1-based line number
+   */
+  line?: (hast: Element, line: number) => Element | void
+  token?: (hast: Element, line: number, col: number) => Element | void
+}
+
 export interface HtmlRendererOptions {
   langId?: string
   fg?: string
   bg?: string
 
-  hastTransform?: (hast: Root) => Root | void
-
-  elements?: ElementsOptions
+  /**
+   * Hast transformers
+   */
+  transforms?: HastTransformers
 
   themeName?: string
 
@@ -270,6 +290,7 @@ export interface HtmlRendererOptions {
    * When specified, `fg` and `bg` will be ignored.
    */
   rootStyle?: string
+
   /**
    * Merge token with only whitespace to the next token,
    * Saving a few extra `<span>`
@@ -279,39 +300,6 @@ export interface HtmlRendererOptions {
   mergeWhitespaces?: boolean
 }
 
-export interface ElementsOptions {
-  pre?: (props: PreElementProps) => string
-  code?: (props: CodeElementProps) => string
-  line?: (props: LineElementProps) => string
-  token?: (props: TokenElementProps) => string
-}
-
-interface ElementProps {
-  children: string
-  [key: string]: unknown
-}
-
-interface PreElementProps extends ElementProps {
-  className: string
-  style: string
-}
-
-interface CodeElementProps extends ElementProps {}
-
-interface LineElementProps extends ElementProps {
-  className: string
-  lines: ThemedToken[][]
-  line: ThemedToken[]
-  index: number
-}
-
-interface TokenElementProps extends ElementProps {
-  style: string
-  tokens: ThemedToken[]
-  token: ThemedToken
-  index: number
-}
-
 export interface ThemedTokenScopeExplanation {
   scopeName: string
   themeMatches: any[]
diff --git a/packages/shikiji/test/hast.test.ts b/packages/shikiji/test/hast.test.ts
@@ -1,6 +1,5 @@
 import { describe, expect, it } from 'vitest'
 import { toHtml } from 'hast-util-to-html'
-import type { Element } from 'hast'
 import { codeToHtml, getHighlighter } from '../src'
 
 describe('should', () => {
@@ -23,12 +22,20 @@ describe('should', () => {
     const code = await codeToHtml('foo\bar', {
       lang: 'js',
       theme: 'vitesse-light',
-      hastTransform(root) {
-        (root.children[0] as Element).properties.class = 'foo'
+      transforms: {
+        line(node, line) {
+          node.properties['data-line'] = line
+        },
+        code(node) {
+          node.properties.class = 'language-js'
+        },
+        token(node, line, col) {
+          node.properties.class = `token:${line}:${col}`
+        },
       },
     })
 
     expect(code)
-      .toMatchInlineSnapshot('"<pre class=\\"foo\\" style=\\"background-color:#ffffff;color:#393a34\\" tabindex=\\"0\\"><code><span class=\\"line\\"><span style=\\"color:#B07D48\\">foo</span><span style=\\"color:#393A34\\"></span><span style=\\"color:#B07D48\\">ar</span></span></code></pre>"')
+      .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>"')
   })
 })
