feat(rehype): support inline codes (#751)

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

Author: fuma-nama

docs/packages/rehype.md

@@ -97,3 +97,39 @@ console.log('3') // highlighted
 console.log('4') // highlighted
 ```
 ````
+
+### Inline Code
+
+You can also highlight inline codes with the `inline` option.
+
+| Option                  | Example          | Description                                                 |
+| ----------------------- | ---------------- | ----------------------------------------------------------- |
+| `false`                 | -                | Disable inline code highlighting (default)                  |
+| `'tailing-curly-colon'` | `let a = 1{:js}` | Highlight with a `{:language}` marker inside the code block |
+
+Enable `inline` on the Rehype plugin:
+
+```ts twoslash
+// @noErrors: true
+import { unified } from 'unified'
+import remarkParse from 'remark-parse'
+import remarkRehype from 'remark-rehype'
+import rehypeStringify from 'rehype-stringify'
+import rehypeShiki from '@shikijs/rehype'
+
+const file = await unified()
+  .use(remarkParse)
+  .use(remarkRehype)
+  .use(rehypeShiki, {
+    inline: 'tailing-curly-colon', // or other options
+    // ...
+  })
+  .use(rehypeStringify)
+  .process(await fs.readFile('./input.md'))
+```
+
+Then you can use inline code in markdown:
+
+```md
+This code `console.log("Hello World"){:js}` will be highlighted.
+```

packages/rehype/src/core.ts

@@ -1,73 +1,15 @@
-import type { CodeOptionsMeta, CodeOptionsThemes, CodeToHastOptions, CodeToHastOptionsCommon, HighlighterGeneric, TransformerOptions } from 'shiki/core'
+import type {
+  CodeToHastOptions,
+  HighlighterGeneric,
+} from 'shiki/core'
 import type { Element, Root } from 'hast'
-import type { BuiltinTheme } from 'shiki'
 import type { Transformer } from 'unified'
 import { toString } from 'hast-util-to-string'
 import { visit } from 'unist-util-visit'
+import { InlineCodeProcessors } from './inline'
+import type { RehypeShikiCoreOptions } from './types'
 
-export interface MapLike<K = any, V = any> {
-  get: (key: K) => V | undefined
-  set: (key: K, value: V) => this
-}
-
-export interface RehypeShikiExtraOptions {
-  /**
-   * Add `language-*` class to code element
-   *
-   * @default false
-   */
-  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
-
-  /**
-   * `mdast-util-to-hast` adds a newline to the end of code blocks
-   *
-   * This option strips that newline from the code block
-   *
-   * @default true
-   * @see https://github.com/syntax-tree/mdast-util-to-hast/blob/f511a93817b131fb73419bf7d24d73a5b8b0f0c2/lib/handlers/code.js#L22
-   */
-  stripEndNewline?: boolean
-
-  /**
-   * Custom meta string parser
-   * Return an object to merge with `meta`
-   */
-  parseMetaString?: (
-    metaString: string,
-    node: Element,
-    tree: Root
-  ) => Record<string, any> | undefined | null
-
-  /**
-   * Custom map to cache transformed codeToHast result
-   *
-   * @default undefined
-   */
-  cache?: MapLike
-
-  /**
-   * Chance to handle the error
-   * If not provided, the error will be thrown
-   */
-  onError?: (error: unknown) => void
-}
-
-export type RehypeShikiCoreOptions =
-  & CodeOptionsThemes<BuiltinTheme>
-  & TransformerOptions
-  & CodeOptionsMeta
-  & RehypeShikiExtraOptions
-  & Omit<CodeToHastOptionsCommon, 'lang'>
+export * from './types'
 
 const languagePrefix = 'language-'
 
@@ -84,86 +26,129 @@ function rehypeShikiFromHighlighter(
     fallbackLanguage,
     onError,
     stripEndNewline = true,
+    inline = false,
     ...rest
   } = options
 
-  return function (tree) {
-    visit(tree, 'element', (node, index, parent) => {
-      if (!parent || index == null || node.tagName !== 'pre')
-        return
-
-      const head = node.children[0]
-
-      if (
-        !head
-        || head.type !== 'element'
-        || head.tagName !== 'code'
-        || !head.properties
-      ) {
-        return
-      }
-
-      const classes = head.properties.className
-      const languageClass = Array.isArray(classes)
-        ? classes.find(
-          d => typeof d === 'string' && d.startsWith(languagePrefix),
-        )
-        : undefined
-
-      let lang = typeof languageClass === 'string' ? languageClass.slice(languagePrefix.length) : defaultLanguage
-
-      if (!lang)
-        return
-
-      if (fallbackLanguage && !langs.includes(lang))
-        lang = fallbackLanguage
-
-      let code = toString(head)
+  /**
+   * Get the determined language of code block (with default language & fallbacks)
+   */
+  function getLanguage(lang = defaultLanguage): string | undefined {
+    if (lang && fallbackLanguage && !langs.includes(lang))
+      return fallbackLanguage
+    return lang
+  }
 
-      if (stripEndNewline && code.endsWith('\n'))
-        code = code.slice(0, -1)
+  function highlight(
+    lang: string,
+    code: string,
+    metaString: string = '',
+    meta: Record<string, unknown> = {},
+  ): Root | undefined {
+    const cacheKey = `${lang}:${metaString}:${code}`
+    const cachedValue = cache?.get(cacheKey)
+
+    if (cachedValue) {
+      return cachedValue
+    }
+
+    const codeOptions: CodeToHastOptions = {
+      ...rest,
+      lang,
+      meta: {
+        ...rest.meta,
+        ...meta,
+        __raw: metaString,
+      },
+    }
+
+    if (addLanguageClass) {
+      // always construct a new array, avoid adding the transformer repeatedly
+      codeOptions.transformers = [
+        ...codeOptions.transformers ?? [],
+        {
+          name: 'rehype-shiki:code-language-class',
+          code(node) {
+            this.addClassToHast(node, `${languagePrefix}${lang}`)
+            return node
+          },
+        },
+      ]
+    }
+
+    if (stripEndNewline && code.endsWith('\n'))
+      code = code.slice(0, -1)
+
+    try {
+      const fragment = highlighter.codeToHast(code, codeOptions)
+      cache?.set(cacheKey, fragment)
+      return fragment
+    }
+    catch (error) {
+      if (onError)
+        onError(error)
+      else
+        throw error
+    }
+  }
 
-      const cachedValue = cache?.get(code)
+  function processPre(tree: Root, node: Element): Root | undefined {
+    const head = node.children[0]
+
+    if (
+      !head
+      || head.type !== 'element'
+      || head.tagName !== 'code'
+      || !head.properties
+    ) {
+      return
+    }
+
+    const classes = head.properties.className
+    const languageClass = Array.isArray(classes)
+      ? classes.find(
+        d => typeof d === 'string' && d.startsWith(languagePrefix),
+      )
+      : undefined
+
+    const lang = getLanguage(
+      typeof languageClass === 'string'
+        ? languageClass.slice(languagePrefix.length)
+        : undefined,
+    )
+
+    if (!lang)
+      return
+
+    const code = toString(head)
+    const metaString = head.data?.meta ?? head.properties.metastring?.toString() ?? ''
+    const meta = parseMetaString?.(metaString, node, tree) || {}
+
+    return highlight(lang, code, metaString, meta)
+  }
 
-      if (cachedValue) {
-        parent.children.splice(index, 1, ...cachedValue)
+  return function (tree) {
+    visit(tree, 'element', (node, index, parent) => {
+      // needed for hast node replacement
+      if (!parent || index == null)
         return
-      }
 
-      const metaString = head.data?.meta ?? head.properties.metastring?.toString() ?? ''
-      const meta = parseMetaString?.(metaString, node, tree) || {}
+      if (node.tagName === 'pre') {
+        const result = processPre(tree, node)
 
-      const codeOptions: CodeToHastOptions = {
-        ...rest,
-        lang,
-        meta: {
-          ...rest.meta,
-          ...meta,
-          __raw: metaString,
-        },
-      }
+        if (result) {
+          parent.children.splice(index, 1, ...result.children)
+        }
 
-      if (addLanguageClass) {
-        codeOptions.transformers ||= []
-        codeOptions.transformers.push({
-          name: 'rehype-shiki:code-language-class',
-          code(node) {
-            this.addClassToHast(node, `${languagePrefix}${lang}`)
-            return node
-          },
-        })
+        // don't look for the `code` node inside
+        return 'skip'
       }
 
-      try {
-        const fragment = highlighter.codeToHast(code, codeOptions)
-        cache?.set(code, fragment.children)
-        parent.children.splice(index, 1, ...fragment.children)
-      }
-      catch (error) {
-        if (onError)
-          onError(error)
-        else
-          throw error
+      if (node.tagName === 'code' && inline) {
+        const result = InlineCodeProcessors[inline]?.({ node, getLanguage, highlight })
+        if (result) {
+          parent.children.splice(index, 1, ...result.children)
+        }
       }
     })
   }

packages/rehype/src/index.ts

@@ -6,7 +6,7 @@ import { bundledLanguages, getSingletonHighlighter } from 'shiki'
 import type { Plugin } from 'unified'
 import type { Root } from 'hast'
 import rehypeShikiFromHighlighter from './core'
-import type { RehypeShikiCoreOptions } from './core'
+import type { RehypeShikiCoreOptions } from './types'
 
 export type RehypeShikiOptions = RehypeShikiCoreOptions
   & {

packages/rehype/src/inline.ts

@@ -0,0 +1,42 @@
+import type { Element, Root } from 'hast'
+import { toString } from 'hast-util-to-string'
+import type { RehypeShikiCoreOptions } from './types'
+
+interface InlineCodeProcessorContext {
+  node: Element
+  getLanguage: (lang?: string) => string | undefined
+  highlight: (
+    lang: string,
+    code: string,
+    metaString?: string,
+    meta?: Record<string, unknown>
+  ) => Root | undefined
+}
+
+type InlineCodeProcessor = (context: InlineCodeProcessorContext) => Root | undefined
+
+type Truthy<T> = T extends false | '' | 0 | null | undefined ? never : T
+
+export const InlineCodeProcessors: Record<Truthy<RehypeShikiCoreOptions['inline']>, InlineCodeProcessor> = {
+  'tailing-curly-colon': ({ node, getLanguage, highlight }) => {
+    const raw = toString(node)
+    const match = raw.match(/(.+)\{:([\w-]+)\}$/)
+    if (!match)
+      return
+    const lang = getLanguage(match[2])
+    if (!lang)
+      return
+
+    const code = match[1] ?? raw
+    const fragment = highlight(lang, code)
+    if (!fragment)
+      return
+
+    const head = fragment.children[0]
+    if (head.type === 'element' && head.tagName === 'pre') {
+      head.tagName = 'span'
+    }
+
+    return fragment
+  },
+}