feat: add codeToTokensWithThemes, fix #7

Co-authored-by: starknt <starknt@users.noreply.github.com>

packages/shikiji/src/core/bundle-factory.ts

@@ -1,4 +1,4 @@
-import type { BundledHighlighterOptions, CodeToHtmlThemesOptions, CodeToHtmlOptions, CodeToThemedTokensOptions, HighlighterCoreOptions, HighlighterGeneric, LanguageInput, MaybeArray, PlainTextLanguage, RequireKeys, ThemeInput } from '../types'
+import type { BundledHighlighterOptions, CodeToHtmlOptions, CodeToHtmlThemesOptions, CodeToThemedTokensOptions, CodeToTokensWithThemesOptions, HighlighterCoreOptions, HighlighterGeneric, LanguageInput, MaybeArray, PlainTextLanguage, RequireKeys, ThemeInput } from '../types'
 import { isPlaintext, toArray } from './utils'
 import { getHighlighterCore } from './core'
 
@@ -71,7 +71,7 @@ export function createdBundledHighlighter<BundledLangs extends string, BundledTh
 export function createSingletonShorthands<L extends string, T extends string >(getHighlighter: GetHighlighterFactory<L, T>) {
   let _shiki: ReturnType<typeof getHighlighter>
 
-  async function getShikiWithThemeLang(options: { theme: MaybeArray<T>; lang: MaybeArray<L | PlainTextLanguage> }) {
+  async function _getHighlighter(options: { theme: MaybeArray<T>; lang: MaybeArray<L | PlainTextLanguage> }) {
     if (!_shiki) {
       _shiki = getHighlighter({
         themes: toArray(options.theme),
@@ -96,7 +96,7 @@ export function createSingletonShorthands<L extends string, T extends string >(g
    * Differences from `shiki.codeToHtml()`, this function is async.
    */
   async function codeToHtml(code: string, options: RequireKeys<CodeToHtmlOptions<L, T>, 'theme' | 'lang'>) {
-    const shiki = await getShikiWithThemeLang(options)
+    const shiki = await _getHighlighter(options)
     return shiki.codeToHtml(code, options)
   }
 
@@ -107,7 +107,7 @@ export function createSingletonShorthands<L extends string, T extends string >(g
    * Differences from `shiki.codeToThemedTokens()`, this function is async.
    */
   async function codeToThemedTokens(code: string, options: RequireKeys<CodeToThemedTokensOptions<L, T>, 'theme' | 'lang'>) {
-    const shiki = await getShikiWithThemeLang(options)
+    const shiki = await _getHighlighter(options)
     return shiki.codeToThemedTokens(code, options)
   }
 
@@ -118,16 +118,31 @@ export function createSingletonShorthands<L extends string, T extends string >(g
    * Differences from `shiki.codeToHtmlThemes()`, this function is async.
    */
   async function codeToHtmlThemes(code: string, options: RequireKeys<CodeToHtmlThemesOptions<L, T>, 'themes' | 'lang'>) {
-    const shiki = await getShikiWithThemeLang({
+    const shiki = await _getHighlighter({
       lang: options.lang,
       theme: Object.values(options.themes).filter(Boolean) as T[],
     })
     return shiki.codeToHtmlThemes(code, options)
   }
 
+  /**
+   * Shorthand for `codeToTokensWithThemes` with auto-loaded theme and language.
+   * A singleton highlighter it maintained internally.
+   *
+   * Differences from `shiki.codeToTokensWithThemes()`, this function is async.
+   */
+  async function codeToTokensWithThemes(code: string, options: RequireKeys<CodeToTokensWithThemesOptions<L, T>, 'themes' | 'lang'>) {
+    const shiki = await _getHighlighter({
+      lang: options.lang,
+      theme: Object.values(options.themes).filter(Boolean) as T[],
+    })
+    return shiki.codeToTokensWithThemes(code, options)
+  }
+
   return {
     codeToHtml,
     codeToHtmlThemes,
     codeToThemedTokens,
+    codeToTokensWithThemes,
   }
 }

packages/shikiji/src/core/core.ts

@@ -1,11 +1,11 @@
-import type { CodeToHtmlOptions, CodeToHtmlThemesOptions, CodeToThemedTokensOptions, HighlighterCoreOptions, HighlighterGeneric, LanguageInput, MaybeGetter, ThemeInput, ThemeRegisteration, ThemedToken } from '../types'
+import type { CodeToHtmlOptions, CodeToHtmlThemesOptions, CodeToThemedTokensOptions, CodeToTokensWithThemesOptions, HighlighterCoreOptions, HighlighterGeneric, LanguageInput, MaybeGetter, ThemeInput, ThemedToken } from '../types'
 import { createOnigScanner, createOnigString, loadWasm } from '../oniguruma'
 import { Registry } from './registry'
 import { Resolver } from './resolver'
 import { tokenizeWithTheme } from './themedTokenizer'
 import { renderToHtml } from './renderer-html'
 import { isPlaintext } from './utils'
-import { renderToHtmlThemes } from './renderer-html-themes'
+import { renderToHtmlThemes, syncThemesTokenization } from './renderer-html-themes'
 
 export type HighlighterCore = HighlighterGeneric<never, never>
 
@@ -85,6 +85,9 @@ export async function getHighlighterCore(options: HighlighterCoreOptions = {}):
     }
   }
 
+  /**
+   * Get highlighted code in HTML.
+   */
   function codeToHtml(code: string, options: CodeToHtmlOptions = {}): string {
     const tokens = codeToThemedTokens(code, {
       ...options,
@@ -99,27 +102,38 @@ export async function getHighlighterCore(options: HighlighterCoreOptions = {}):
     })
   }
 
+  /**
+   * Get tokens with multiple themes, with synced
+   */
+  function codeToTokensWithThemes(code: string, options: CodeToTokensWithThemesOptions) {
+    const themes = Object.entries(options.themes)
+      .filter(i => i[1]) as [string, string][]
+
+    const tokens = syncThemesTokenization(...themes.map(t => codeToThemedTokens(code, {
+      ...options,
+      theme: t[1],
+      includeExplanation: false,
+    })))
+
+    return themes.map(([color, theme], idx) => [
+      color,
+      theme,
+      tokens[idx],
+    ] as [string, string, ThemedToken[][]])
+  }
+
   function codeToHtmlThemes(code: string, options: CodeToHtmlThemesOptions): string {
     const {
-      defaultColor = 'light', cssVariablePrefix = '--shiki-',
+      defaultColor = 'light',
+      cssVariablePrefix = '--shiki-',
     } = options
 
-    const themes = Object.entries(options.themes)
-      .filter(i => i[1])
+    const tokens = codeToTokensWithThemes(code, options)
       .sort(a => a[0] === defaultColor ? -1 : 1)
 
-    const tokens = themes.map(([color, theme]) => [
-      color,
-      codeToThemedTokens(code, {
-        ...options,
-        theme,
-        includeExplanation: false,
-      }),
-      getTheme(theme),
-    ] as [string, ThemedToken[][], ThemeRegisteration])
-
     return renderToHtmlThemes(
       tokens,
+      tokens.map(i => getTheme(i[1])),
       cssVariablePrefix,
       defaultColor !== false,
       options,
@@ -132,15 +146,15 @@ export async function getHighlighterCore(options: HighlighterCoreOptions = {}):
 
   async function loadTheme(...themes: ThemeInput[]) {
     await Promise.all(
-      themes.map(async theme => _registry.loadTheme(await normalizeGetter(theme)),
-      ),
+      themes.map(async theme => _registry.loadTheme(await normalizeGetter(theme))),
     )
   }
 
   return {
-    codeToThemedTokens,
     codeToHtml,
     codeToHtmlThemes,
+    codeToThemedTokens,
+    codeToTokensWithThemes,
     loadLanguage,
     loadTheme,
     getLoadedThemes: () => _registry.getLoadedThemes(),

packages/shikiji/src/core/renderer-html-themes.ts

@@ -2,19 +2,19 @@ import type { HtmlRendererOptions, ThemeRegisteration, ThemedToken } from '../ty
 import { renderToHtml } from './renderer-html'
 
 /**
- * Break tokens from multiple themes into same length.
+ * Break tokens from multiple themes into same tokenization.
  *
  * For example, given two themes that tokenize `console.log("hello")` as:
  *
  * - `console . log (" hello ")` (6 tokens)
  * - `console .log ( "hello" )` (5 tokens)
  *
- * This function break both themes into same tokenization, so later the can be rendered in pairs.
+ * This function will return:
  *
  * - `console . log ( " hello " )` (8 tokens)
  * - `console . log ( " hello " )` (8 tokens)
  */
-export function _syncThemedTokens(...themes: ThemedToken[][][]) {
+export function syncThemesTokenization(...themes: ThemedToken[][][]) {
   const outThemes = themes.map<ThemedToken[][]>(() => [])
   const count = themes.length
 
@@ -55,35 +55,35 @@ export function _syncThemedTokens(...themes: ThemedToken[][][]) {
 }
 
 export function renderToHtmlThemes(
-  themes: [string, ThemedToken[][], ThemeRegisteration][],
+  themeTokens: [string, string, ThemedToken[][]][],
+  themeRegs: ThemeRegisteration[],
   cssVariablePrefix = '--shiki-',
   defaultColor = true,
   options: HtmlRendererOptions = {},
 ) {
-  const synced = _syncThemedTokens(...themes.map(t => t[1]))
-
+  const themeMap = themeTokens.map(t => t[2])
   const merged: ThemedToken[][] = []
-  for (let i = 0; i < synced[0].length; i++) {
-    const lines = synced.map(t => t[i])
+  for (let i = 0; i < themeMap[0].length; i++) {
+    const lineMap = themeMap.map(t => t[i])
     const lineout: any[] = []
     merged.push(lineout)
-    for (let j = 0; j < lines[0].length; j++) {
-      const tokens = lines.map(t => t[j])
-      const colors = tokens.map((t, idx) => `${idx === 0 && defaultColor ? '' : `${cssVariablePrefix + themes[idx][0]}:`}${t.color || 'inherit'}`).join(';')
+    for (let j = 0; j < lineMap[0].length; j++) {
+      const tokenMap = lineMap.map(t => t[j])
+      const colors = tokenMap.map((t, idx) => `${idx === 0 && defaultColor ? '' : `${cssVariablePrefix + themeTokens[idx][0]}:`}${t.color || 'inherit'}`).join(';')
       lineout.push({
-        ...tokens[0],
+        ...tokenMap[0],
         color: colors,
         htmlStyle: defaultColor ? undefined : colors,
       })
     }
   }
 
-  const fg = options.fg || themes.map((t, idx) => (idx === 0 && defaultColor ? '' : `${cssVariablePrefix + t[0]}:`) + t[2].fg).join(';')
-  const bg = options.bg || themes.map((t, idx) => (idx === 0 && defaultColor ? '' : `${cssVariablePrefix + t[0]}-bg:`) + t[2].bg).join(';')
+  const fg = options.fg || themeTokens.map((t, idx) => (idx === 0 && defaultColor ? '' : `${cssVariablePrefix + t[0]}:`) + themeRegs[idx].fg).join(';')
+  const bg = options.bg || themeTokens.map((t, idx) => (idx === 0 && defaultColor ? '' : `${cssVariablePrefix + t[0]}-bg:`) + themeRegs[idx].bg).join(';')
   return renderToHtml(merged, {
     fg,
     bg,
-    themeName: `shiki-themes ${themes.map(t => t[2].name).join(' ')}`,
+    themeName: `shiki-themes ${themeRegs.map(t => t.name).join(' ')}`,
     rootStyle: defaultColor ? undefined : [fg, bg].join(';'),
     ...options,
   })

packages/shikiji/src/core/themedTokenizer.ts

@@ -5,7 +5,7 @@
 
 import type { IGrammar, IRawTheme } from 'vscode-textmate'
 import { INITIAL } from 'vscode-textmate'
-import type { ThemedToken, ThemedTokenExplanation, ThemedTokenScopeExplanation } from '../types'
+import type { ThemedToken, ThemedTokenScopeExplanation } from '../types'
 import type { FontStyle } from './stackElementMetadata'
 import { StackElementMetadata } from './stackElementMetadata'
 
@@ -58,8 +58,14 @@ export function tokenizeWithTheme(
       const foregroundColor = colorMap[foreground]
       const fontStyle: FontStyle = StackElementMetadata.getFontStyle(metadata)
 
-      const explanation: ThemedTokenExplanation[] = []
+      const token: ThemedToken = {
+        content: line.substring(startIndex, nextStartIndex),
+        color: foregroundColor,
+        fontStyle,
+      }
+
       if (options.includeExplanation) {
+        token.explanation = []
         let offset = 0
         while (startIndex + offset < nextStartIndex) {
           const tokenWithScopes = tokensWithScopes![tokensWithScopesIndex!]
@@ -69,7 +75,7 @@ export function tokenizeWithTheme(
             tokenWithScopes.endIndex,
           )
           offset += tokenWithScopesText.length
-          explanation.push({
+          token.explanation.push({
             content: tokenWithScopesText,
             scopes: explainThemeScopes(theme, tokenWithScopes.scopes),
           })
@@ -78,12 +84,7 @@ export function tokenizeWithTheme(
         }
       }
 
-      actual.push({
-        content: line.substring(startIndex, nextStartIndex),
-        color: foregroundColor,
-        fontStyle,
-        explanation,
-      })
+      actual.push(token)
     }
     final.push(actual)
     actual = []

packages/shikiji/src/index.ts

@@ -24,6 +24,7 @@ export const {
   codeToHtml,
   codeToHtmlThemes,
   codeToThemedTokens,
+  codeToTokensWithThemes,
 } = /* @__PURE__ */ createSingletonShorthands<
   BuiltinLanguages,
   BuiltinThemes

packages/shikiji/src/types.ts

@@ -42,6 +42,10 @@ export interface HighlighterGeneric<BundledLangKeys extends string, BundledTheme
     code: string,
     options: CodeToThemedTokensOptions<ResolveBundleKey<BundledLangKeys>, ResolveBundleKey<BundledThemeKeys>>
   ): ThemedToken[][]
+  codeToTokensWithThemes(
+    code: string,
+    options: CodeToTokensWithThemesOptions<ResolveBundleKey<BundledLangKeys>, ResolveBundleKey<BundledThemeKeys>>
+  ): [color: string, theme: string, tokens: ThemedToken[][]][]
   loadTheme(...themes: (ThemeInput | BundledThemeKeys)[]): Promise<void>
   loadLanguage(...langs: (LanguageInput | BundledLangKeys | PlainTextLanguage)[]): Promise<void>
   getLoadedLanguages(): string[]
@@ -84,13 +88,26 @@ export interface LanguageRegistration extends IRawGrammar {
   unbalancedBracketSelectors?: string[]
 }
 
-export interface CodeToHtmlOptions<Languages = string, Themes = string> {
+export interface CodeToThemedTokensOptions<Languages = string, Themes = string> {
   lang?: Languages | PlainTextLanguage
   theme?: Themes
+  /**
+   * Include explanation of why a token is given a color.
+   *
+   * @default true
+   */
+  includeExplanation?: boolean
+}
+
+export interface CodeToHtmlBasicOptions {
   lineOptions?: LineOption[]
 }
 
-export interface CodeToHtmlThemesOptions<Languages = string, Themes = string> {
+export interface CodeToHtmlOptions<Languages = string, Themes = string>
+  extends Omit<CodeToThemedTokensOptions<Languages, Themes>, 'includeExplanation'>, CodeToHtmlBasicOptions {
+}
+
+export interface CodeToTokensWithThemesOptions<Languages = string, Themes = string> {
   lang?: Languages | PlainTextLanguage
 
   /**
@@ -112,6 +129,10 @@ export interface CodeToHtmlThemesOptions<Languages = string, Themes = string> {
     light: Themes
     dark: Themes
   } & Partial<Record<string, Themes>>
+}
+
+export interface CodeToHtmlThemesOptions<Languages = string, Themes = string>
+  extends CodeToTokensWithThemesOptions<Languages, Themes>, CodeToHtmlBasicOptions {
 
   /**
    * The default theme applied to the code (via inline `color` style).
@@ -141,19 +162,6 @@ export interface CodeToHtmlThemesOptions<Languages = string, Themes = string> {
    * @default '--shiki-'
    */
   cssVariablePrefix?: string
-
-  lineOptions?: LineOption[]
-}
-
-export interface CodeToThemedTokensOptions<Languages = string, Themes = string> {
-  lang?: Languages | PlainTextLanguage
-  theme?: Themes
-  /**
-   * Include explanation of why a token is given a color.
-   *
-   * @default true
-   */
-  includeExplanation?: boolean
 }
 
 export interface LineOption {