From d12e23ddf69480418078ff39846c99ecf2e1eb1b Mon Sep 17 00:00:00 2001
From: Anthony Fu <anthonyfu117@hotmail.com>
Date: Sat, 16 Dec 2023 06:33:50 +0100
Subject: [PATCH] feat: expose `shikijiSetup` hook (#3344)

---
 src/node/markdown/markdown.ts          | 18 +++++++-----------
 src/node/markdown/plugins/highlight.ts | 22 +++++++++++++---------
 2 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/src/node/markdown/markdown.ts b/src/node/markdown/markdown.ts
index 67c631a1c567..c549901d66aa 100644
--- a/src/node/markdown/markdown.ts
+++ b/src/node/markdown/markdown.ts
@@ -31,7 +31,8 @@ import type {
   ThemeRegistration,
   BuiltinTheme,
   LanguageInput,
-  ShikijiTransformer
+  ShikijiTransformer,
+  Highlighter
 } from 'shikiji'
 
 export type { Header } from '../shared'
@@ -103,6 +104,10 @@ export interface MarkdownOptions extends MarkdownIt.Options {
    * @see https://github.com/antfu/shikiji#hast-transformers
    */
   codeTransformers?: ShikijiTransformer[]
+  /**
+   * Setup Shikiji instance
+   */
+  shikijiSetup?: (shikiji: Highlighter) => void | Promise<void>
 
   /* ==================== Markdown It Plugins ==================== */
 
@@ -176,16 +181,7 @@ export const createMarkdownRenderer = async (
   const md = MarkdownIt({
     html: true,
     linkify: true,
-    highlight:
-      options.highlight ||
-      (await highlight(
-        theme,
-        options.languages,
-        options.defaultHighlightLang,
-        logger,
-        options.codeTransformers,
-        options.languageAlias
-      )),
+    highlight: options.highlight || (await highlight(theme, options, logger)),
     ...options
   })
 
diff --git a/src/node/markdown/plugins/highlight.ts b/src/node/markdown/plugins/highlight.ts
index 17b820a14d10..84d999e6a3be 100644
--- a/src/node/markdown/plugins/highlight.ts
+++ b/src/node/markdown/plugins/highlight.ts
@@ -1,6 +1,6 @@
 import { customAlphabet } from 'nanoid'
 import c from 'picocolors'
-import type { LanguageInput, ShikijiTransformer } from 'shikiji'
+import type { ShikijiTransformer } from 'shikiji'
 import {
   bundledLanguages,
   getHighlighter,
@@ -9,7 +9,7 @@ import {
   isSpecialLang
 } from 'shikiji'
 import type { Logger } from 'vite'
-import type { ThemeOptions } from '../markdown'
+import type { MarkdownOptions, ThemeOptions } from '../markdown'
 import {
   transformerCompactLineOptions,
   transformerNotationDiff,
@@ -55,21 +55,25 @@ const attrsToLines = (attrs: string): TransformerCompactLineOption[] => {
 
 export async function highlight(
   theme: ThemeOptions,
-  languages: LanguageInput[] = [],
-  defaultLang: string = '',
-  logger: Pick<Logger, 'warn'> = console,
-  userTransformers: ShikijiTransformer[] = [],
-  languageAlias: Record<string, string> = {}
+  options: MarkdownOptions,
+  logger: Pick<Logger, 'warn'> = console
 ): Promise<(str: string, lang: string, attrs: string) => string> {
+  const {
+    defaultHighlightLang: defaultLang = '',
+    codeTransformers: userTransformers = []
+  } = options
+
   const highlighter = await getHighlighter({
     themes:
       typeof theme === 'string' || 'name' in theme
         ? [theme]
         : [theme.light, theme.dark],
-    langs: [...Object.keys(bundledLanguages), ...languages],
-    langAlias: languageAlias
+    langs: [...Object.keys(bundledLanguages), ...(options.languages || [])],
+    langAlias: options.languageAlias
   })
 
+  await options?.shikijiSetup?.(highlighter)
+
   const transformers: ShikijiTransformer[] = [
     transformerNotationDiff(),
     transformerNotationFocus({