feat: add addClassToHast to transformer context (#573)

shikijs · Jan 31, 2024 · 6dfcd98 · 6dfcd98
1 parent d123b4b
commit 6dfcd98
Show file tree

Hide file tree

Showing 14 changed files with 34 additions and 29 deletions.
diff --git a/docs/.vitepress/store/playground.ts b/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 || ''
               },
             },

diff --git a/docs/guide/transformers.md b/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}`

diff --git a/packages/core/src/code-to-hast.ts b/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
     },

diff --git a/packages/core/src/types.ts b/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 {

diff --git a/packages/core/src/utils.ts b/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
 }
 
 /**

diff --git a/packages/rehype/src/core.ts b/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
           },
         })

diff --git a/packages/transformers/src/shared/highlight-word.ts b/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[]

diff --git a/packages/transformers/src/transformers/compact-line-options.ts b/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
     },
   }

diff --git a/packages/transformers/src/transformers/notation-highlight-word.ts b/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

diff --git a/packages/transformers/src/transformers/notation-map.ts b/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
     },
   )

diff --git a/packages/transformers/src/transformers/render-whitespace.ts b/packages/transformers/src/transformers/render-whitespace.ts
@@ -1,5 +1,4 @@
 import type { ShikiTransformer } from 'shiki'
-import { addClassToHast } from 'shiki'
 import type { Element } from 'hast'
 import { splitSpaces } from '../shared/utils'
 
@@ -81,7 +80,7 @@ export function transformerRenderWhitespace(
             }
             clone.children = [{ type: 'text', value: part }]
             if (keys.includes(part)) {
-              addClassToHast(clone, classMap[part])
+              this.addClassToHast(clone, classMap[part])
               delete clone.properties.style
             }
             return clone

diff --git a/packages/transformers/src/transformers/transformer-meta-highlight-word.ts b/packages/transformers/src/transformers/transformer-meta-highlight-word.ts
@@ -41,7 +41,7 @@ export function transformerMetaWordHighlight(
       const words = parseMetaHighlightWords(this.options.meta.__raw)
 
       for (const word of words)
-        highlightWordInLine(node, null, word, className)
+        highlightWordInLine.call(this, node, null, word, className)
 
       return node
     },

diff --git a/packages/transformers/src/transformers/transformer-meta-highlight.ts b/packages/transformers/src/transformers/transformer-meta-highlight.ts
@@ -1,4 +1,4 @@
-import { type ShikiTransformer, addClassToHast } from 'shiki'
+import type { ShikiTransformer } from 'shiki'
 
 export function parseMetaHighlightString(meta: string) {
   if (!meta)
@@ -46,7 +46,7 @@ export function transformerMetaHighlight(
       ;(this.meta as any)[symbol] ||= parseMetaHighlightString(this.options.meta.__raw)
       const lines: number[] = (this.meta as any)[symbol] || []
       if (lines.includes(line))
-        addClassToHast(node, className)
+        this.addClassToHast(node, className)
       return node
     },
   }

diff --git a/packages/twoslash/src/core.ts b/packages/twoslash/src/core.ts
@@ -6,7 +6,7 @@ import type { TwoslashExecuteOptions, TwoslashReturn } from 'twoslash'
 import type { ShikiTransformer } from '@shikijs/core'
 import type { Element, ElementContent, Text } from 'hast'
 
-import { addClassToHast, splitToken } from '@shikijs/core'
+import { splitToken } from '@shikijs/core'
 import type { TransformerTwoslashOptions, TwoslashRenderer } from './types'
 
 export * from './types'
@@ -86,7 +86,7 @@ export function createTransformerFactory(
       },
       pre(pre) {
         if (this.meta.twoslash)
-          addClassToHast(pre, 'twoslash lsp')
+          this.addClassToHast(pre, 'twoslash lsp')
       },
       code(codeEl) {
         const twoslash = this.meta.twoslash