feat: add addClassToHast to transformer context

Author: antfu

docs/.vitepress/store/playground.ts

@@ -39,7 +39,7 @@ export const usePlayground = defineStore('playground', () => {
   }
 
   ;(async () => {
-    const { getHighlighter, addClassToHast } = await import('shiki')
+    const { getHighlighter } = await import('shiki')
     const { bundledLanguagesInfo: bundleFull } = await import('shiki/bundle/full')
     const { bundledLanguagesInfo: bundleWeb } = await import('shiki/bundle/web')
     const { bundledThemesInfo } = await import('shiki/themes')
@@ -93,7 +93,7 @@ export const usePlayground = defineStore('playground', () => {
           transformers: [
             {
               pre(node) {
-                addClassToHast(node, 'vp-code')
+                this.addClassToHast(node, 'vp-code')
                 preStyle.value = node.properties?.style as string || ''
               },
             },

docs/guide/transformers.md

@@ -5,20 +5,20 @@ Shiki uses [`hast`](https://github.com/syntax-tree/hast), a AST format for HTML,
 You can provide your own `transformers` 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:
 
 ```ts twoslash
-import { addClassToHast, codeToHtml } from 'shiki'
+import { codeToHtml } from 'shiki'
 
 const code = await codeToHtml('foo\bar', {
   lang: 'js',
   theme: 'vitesse-light',
   transformers: [
     {
       code(node) {
-        addClassToHast(node, 'language-js')
+        this.addClassToHast(node, 'language-js')
       },
       line(node, line) {
         node.properties['data-line'] = line
         if ([1, 3, 4].includes(line))
-          addClassToHast(node, 'highlight')
+          this.addClassToHast(node, 'highlight')
       },
       span(node, line, col) {
         node.properties['data-token'] = `token:${line}:${col}`

packages/core/src/code-to-hast.ts

@@ -9,7 +9,7 @@ import type {
 } from './types'
 import { FontStyle } from './types'
 import { codeToTokens } from './code-to-tokens'
-import { getTokenStyleObject, stringifyTokenStyle } from './utils'
+import { addClassToHast, getTokenStyleObject, stringifyTokenStyle } from './utils'
 
 export function codeToHast(
   internal: ShikiInternal,
@@ -102,6 +102,7 @@ export function tokensToHast(
 
   const context: ShikiTransformerContext = {
     ...transformerContext,
+    addClassToHast,
     get tokens() {
       return tokens
     },

packages/core/src/types.ts

@@ -488,6 +488,13 @@ export interface ShikiTransformerContext extends ShikiTransformerContextCommon {
   readonly pre: Element
   readonly code: Element
   readonly lines: Element[]
+
+  /**
+   * Utility to append class to a hast node
+   *
+   * If the `property.class` is a string, it will be splitted by space and converted to an array.
+   */
+  addClassToHast: (hast: Element, className: string | string[]) => Element
 }
 
 export interface ShikiTransformer {

packages/core/src/utils.ts

@@ -56,7 +56,7 @@ export function isSpecialTheme(theme: string | ThemeInput | null | undefined): t
  */
 export function addClassToHast(node: Element, className: string | string[]) {
   if (!className)
-    return
+    return node
   node.properties ||= {}
   node.properties.class ||= []
   if (typeof node.properties.class === 'string')
@@ -69,6 +69,7 @@ export function addClassToHast(node: Element, className: string | string[]) {
     if (c && !node.properties.class.includes(c))
       node.properties.class.push(c)
   }
+  return node
 }
 
 /**

packages/rehype/src/core.ts

@@ -1,4 +1,3 @@
-import { addClassToHast } from 'shiki/core'
 import type { CodeOptionsMeta, CodeOptionsThemes, CodeToHastOptions, HighlighterGeneric, TransformerOptions } from 'shiki/core'
 import type { Element, Root } from 'hast'
 import type { BuiltinTheme } from 'shiki'
@@ -135,7 +134,7 @@ const rehypeShikiFromHighlighter: Plugin<[HighlighterGeneric<any, any>, RehypeSh
         codeOptions.transformers.push({
           name: 'rehype-shiki:code-language-class',
           code(node) {
-            addClassToHast(node, language)
+            this.addClassToHast(node, language)
             return node
           },
         })

packages/transformers/src/shared/highlight-word.ts

@@ -1,12 +1,12 @@
 import type { Element, ElementContent, Text } from 'hast'
-import { addClassToHast } from 'shiki/core'
+import type { ShikiTransformerContext } from 'shiki/core'
 
-export function highlightWordInLine(line: Element, ignoredElement: Element | null, word: string, className: string): void {
+export function highlightWordInLine(this: ShikiTransformerContext, line: Element, ignoredElement: Element | null, word: string, className: string): void {
   const content = getTextContent(line)
   let index = content.indexOf(word)
 
   while (index !== -1) {
-    highlightRange(line.children, ignoredElement, index, word.length, className)
+    highlightRange.call(this, line.children, ignoredElement, index, word.length, className)
     index = content.indexOf(word, index + 1)
   }
 }
@@ -25,7 +25,7 @@ function getTextContent(element: ElementContent): string {
  * @param index highlight beginning index
  * @param len highlight length
  */
-function highlightRange(elements: ElementContent[], ignoredElement: Element | null, index: number, len: number, className: string) {
+function highlightRange(this: ShikiTransformerContext, elements: ElementContent[], ignoredElement: Element | null, index: number, len: number, className: string) {
   let currentIdx = 0
 
   for (let i = 0; i < elements.length; i++) {
@@ -45,7 +45,7 @@ function highlightRange(elements: ElementContent[], ignoredElement: Element | nu
         continue
 
       const separated = separateToken(element, textNode, start, length)
-      addClassToHast(separated[1], className)
+      this.addClassToHast(separated[1], className)
 
       // insert
       const output = separated.filter(Boolean) as Element[]

packages/transformers/src/transformers/compact-line-options.ts

@@ -1,5 +1,4 @@
 import type { ShikiTransformer } from 'shiki'
-import { addClassToHast } from 'shiki'
 
 export interface TransformerCompactLineOption {
   /**
@@ -20,7 +19,7 @@ export function transformerCompactLineOptions(
     line(node, line) {
       const lineOption = lineOptions.find(o => o.line === line)
       if (lineOption?.classes)
-        addClassToHast(node, lineOption.classes)
+        this.addClassToHast(node, lineOption.classes)
       return node
     },
   }

packages/transformers/src/transformers/notation-highlight-word.ts

@@ -1,4 +1,4 @@
-import { type ShikiTransformer, addClassToHast } from 'shiki'
+import type { ShikiTransformer } from 'shiki'
 import { createCommentNotationTransformer } from '../utils'
 import { highlightWordInLine } from '../shared/highlight-word'
 
@@ -34,10 +34,10 @@ export function transformerNotationWordHighlight(
       lines
         // Don't include the comment itself
         .slice(index + 1, index + 1 + lineNum)
-        .forEach(line => highlightWordInLine(line, comment, word, classActiveWord))
+        .forEach(line => highlightWordInLine.call(this, line, comment, word, classActiveWord))
 
       if (classActivePre)
-        addClassToHast(this.pre, classActivePre)
+        this.addClassToHast(this.pre, classActivePre)
       return true
     },
     true, // remove empty lines

packages/transformers/src/transformers/notation-map.ts

@@ -1,5 +1,4 @@
 import type { ShikiTransformer } from 'shiki'
-import { addClassToHast } from 'shiki'
 import { createCommentNotationTransformer } from '../utils'
 
 export interface TransformerNotationMapOptions {
@@ -31,10 +30,10 @@ export function transformerNotationMap(
       lines
         .slice(index, index + lineNum)
         .forEach((line) => {
-          addClassToHast(line, classMap[match])
+          this.addClassToHast(line, classMap[match])
         })
       if (classActivePre)
-        addClassToHast(this.pre, classActivePre)
+        this.addClassToHast(this.pre, classActivePre)
       return true
     },
   )