feat(rehype): Support default language option (#692)

Co-authored-by: Anthony Fu <[email protected]>

packages/rehype/src/core.ts

@@ -1,7 +1,7 @@
 import type { CodeOptionsMeta, CodeOptionsThemes, CodeToHastOptions, HighlighterGeneric, TransformerOptions } from 'shiki/core'
 import type { Element, Root } from 'hast'
 import type { BuiltinTheme } from 'shiki'
-import type { Plugin } from 'unified'
+import type { Transformer } from 'unified'
 import { toString } from 'hast-util-to-string'
 import { visit } from 'unist-util-visit'
 
@@ -18,6 +18,16 @@ export interface RehypeShikiExtraOptions {
    */
   addLanguageClass?: boolean
 
+  /**
+   * The default language to use when is not specified
+   */
+  defaultLanguage?: string
+
+  /**
+   * The fallback language to use when specified language is not loaded
+   */
+  fallbackLanguage?: string
+
   /**
    * Custom meta string parser
    * Return an object to merge with `meta`
@@ -57,19 +67,23 @@ declare module 'hast' {
   }
 }
 
-const rehypeShikiFromHighlighter: Plugin<[HighlighterGeneric<any, any>, RehypeShikiCoreOptions], Root> = function (
-  highlighter,
-  options,
-) {
+const languagePrefix = 'language-'
+
+function rehypeShikiFromHighlighter(
+  highlighter: HighlighterGeneric<any, any>,
+  options: RehypeShikiCoreOptions,
+): Transformer<Root, Root> {
+  const langs = highlighter.getLoadedLanguages()
   const {
     addLanguageClass = false,
     parseMetaString,
     cache,
+    defaultLanguage,
+    fallbackLanguage,
+    onError,
     ...rest
   } = options
 
-  const prefix = 'language-'
-
   return function (tree) {
     visit(tree, 'element', (node, index, parent) => {
       if (!parent || index == null || node.tagName !== 'pre')
@@ -87,19 +101,21 @@ const rehypeShikiFromHighlighter: Plugin<[HighlighterGeneric<any, any>, RehypeSh
       }
 
       const classes = head.properties.className
+      const languageClass = Array.isArray(classes)
+        ? classes.find(
+          d => typeof d === 'string' && d.startsWith(languagePrefix),
+        )
+        : undefined
 
-      if (!Array.isArray(classes))
-        return
-
-      const language = classes.find(
-        d => typeof d === 'string' && d.startsWith(prefix),
-      )
+      let lang = typeof languageClass === 'string' ? languageClass.slice(languagePrefix.length) : defaultLanguage
 
-      if (typeof language !== 'string')
+      if (!lang)
         return
 
-      const code = toString(head as any)
+      if (fallbackLanguage && !langs.includes(lang))
+        lang = fallbackLanguage
 
+      const code = toString(head)
       const cachedValue = cache?.get(code)
 
       if (cachedValue) {
@@ -112,7 +128,7 @@ const rehypeShikiFromHighlighter: Plugin<[HighlighterGeneric<any, any>, RehypeSh
 
       const codeOptions: CodeToHastOptions = {
         ...rest,
-        lang: language.slice(prefix.length),
+        lang,
         meta: {
           ...rest.meta,
           ...meta,
@@ -125,7 +141,7 @@ const rehypeShikiFromHighlighter: Plugin<[HighlighterGeneric<any, any>, RehypeSh
         codeOptions.transformers.push({
           name: 'rehype-shiki:code-language-class',
           code(node) {
-            this.addClassToHast(node, language)
+            this.addClassToHast(node, `${languagePrefix}${lang}`)
             return node
           },
         })
@@ -137,8 +153,8 @@ const rehypeShikiFromHighlighter: Plugin<[HighlighterGeneric<any, any>, RehypeSh
         parent.children.splice(index, 1, ...fragment.children)
       }
       catch (error) {
-        if (options.onError)
-          options.onError(error)
+        if (onError)
+          onError(error)
         else
           throw error
       }

packages/rehype/src/index.ts

@@ -17,24 +17,22 @@ export type RehypeShikiOptions = RehypeShikiCoreOptions
   }
 
 const rehypeShiki: Plugin<[RehypeShikiOptions], Root> = function (
-  options = {} as any,
+  options = {} as RehypeShikiOptions,
 ) {
   const themeNames = ('themes' in options ? Object.values(options.themes) : [options.theme]).filter(Boolean) as BuiltinTheme[]
   const langs = options.langs || Object.keys(bundledLanguages)
 
-  // eslint-disable-next-line ts/no-this-alias
-  const ctx = this
-  let promise: Promise<any>
+  let getHandler: Promise<any>
 
-  return async function (tree) {
-    if (!promise) {
-      promise = getHighlighter({
+  return async (tree) => {
+    if (!getHandler) {
+      getHandler = getHighlighter({
         themes: themeNames,
         langs,
       })
-        .then(highlighter => rehypeShikiFromHighlighter.call(ctx, highlighter, options))
+        .then(highlighter => rehypeShikiFromHighlighter.call(this, highlighter, options))
     }
-    const handler = await promise
+    const handler = await getHandler
     return handler!(tree) as Root
   }
 }

packages/rehype/test/core.test.ts

@@ -27,6 +27,7 @@ it('run', async () => {
     .use(remarkRehype)
     .use(rehypeShikiFromHighlighter, highlighter, {
       theme: 'vitesse-light',
+      defaultLanguage: 'text',
       transformers: [
         transformerMetaHighlight(),
       ],
@@ -63,6 +64,7 @@ it('run with rehype-raw', async () => {
     .use(rehypeRaw)
     .use(rehypeShikiFromHighlighter, highlighter, {
       theme: 'vitesse-light',
+      defaultLanguage: 'text',
       transformers: [
         transformerMetaHighlight(),
       ],

packages/rehype/test/index.test.ts

@@ -13,6 +13,7 @@ it('run', async () => {
     .use(remarkRehype)
     .use(rehypeShiki, {
       theme: 'vitesse-light',
+      defaultLanguage: 'text',
       transformers: [
         transformerMetaWordHighlight(),
         transformerMetaHighlight(),