feat!: improve return type of codeToTokensWithThemes, close #37

antfu · antfu · commit 3acf1bf20609 · 2023-12-07T21:58:12.000+01:00
diff --git a/packages/shikiji/src/core/renderer-hast.ts b/packages/shikiji/src/core/renderer-hast.ts
@@ -1,5 +1,14 @@
 import type { Element, Root, Text } from 'hast'
-import type { CodeToHastOptions, HtmlRendererOptions, ShikiContext, ShikijiTransformerContext, ThemedToken } from '../types'
+import type {
+  CodeToHastOptions,
+  HtmlRendererOptions,
+  ShikiContext,
+  ShikijiTransformerContext,
+  ThemedToken,
+  ThemedTokenWithVariants,
+  TokenStyles,
+} from '../types'
+import css from '../assets/langs/css'
 import { codeToThemedTokens } from './tokenizer'
 import { FontStyle } from './stackElementMetadata'
 import { codeToTokensWithThemes } from './renderer-html-themes'
@@ -22,7 +31,9 @@ export function codeToHast(
     } = options
 
     const themes = Object.entries(options.themes)
-      .filter(i => i[1]) as [string, string][]
+      .filter(i => i[1])
+      .map(i => ({ color: i[0], theme: i[1]! }))
+      .sort((a, b) => a.color === defaultColor ? -1 : b.color === defaultColor ? 1 : 0)
 
     if (themes.length === 0)
       throw new Error('[shikiji] `themes` option must not be empty')
@@ -32,55 +43,18 @@ export function codeToHast(
       code,
       options,
     )
-      .sort(a => a[0] === defaultColor ? -1 : 1)
 
-    if (defaultColor && !themeTokens.find(t => t[0] === defaultColor))
+    if (defaultColor && !themes.find(t => t.color === defaultColor))
       throw new Error(`[shikiji] \`themes\` option must contain the defaultColor key \`${defaultColor}\``)
 
-    const themeRegs = themeTokens.map(t => context.getTheme(t[1]))
-    const themeMap = themeTokens.map(t => t[2])
-    tokens = []
-
-    for (let i = 0; i < themeMap[0].length; i++) {
-      const lineMap = themeMap.map(t => t[i])
-      const lineout: any[] = []
-      tokens.push(lineout)
-      for (let j = 0; j < lineMap[0].length; j++) {
-        const tokenMap = lineMap.map(t => t[j])
-        const tokenStyles = tokenMap.map(t => getTokenStyles(t))
-
-        // Get all style keys, for themes that missing some style, we put `inherit` to override as needed
-        const styleKeys = new Set(tokenStyles.flatMap(t => Object.keys(t)))
-        const mergedStyles = tokenStyles.reduce((acc, cur, idx) => {
-          for (const key of styleKeys) {
-            const value = cur[key] || 'inherit'
-
-            if (idx === 0 && defaultColor) {
-              acc[key] = value
-            }
-            else {
-              const varKey = cssVariablePrefix + themeTokens[idx][0] + (key === 'color' ? '' : `-${key}`)
-              if (acc[key])
-                acc[key] += `;${varKey}:${value}`
-              else
-                acc[key] = `${varKey}:${value}`
-            }
-          }
-          return acc
-        }, {} as Record<string, string>)
-
-        lineout.push({
-          ...tokenMap[0],
-          color: '',
-          htmlStyle: defaultColor
-            ? stringifyTokenStyle(mergedStyles)
-            : Object.values(mergedStyles).join(';'),
-        })
-      }
-    }
+    const themeRegs = themes.map(t => context.getTheme(t.theme))
 
-    fg = themeTokens.map((t, idx) => (idx === 0 && defaultColor ? '' : `${cssVariablePrefix + t[0]}:`) + themeRegs[idx].fg).join(';')
-    bg = themeTokens.map((t, idx) => (idx === 0 && defaultColor ? '' : `${cssVariablePrefix + t[0]}-bg:`) + themeRegs[idx].bg).join(';')
+    const themesOrder = themes.map(t => t.color)
+    tokens = themeTokens
+      .map(line => line.map(token => flattenToken(token, themesOrder, cssVariablePrefix, defaultColor)))
+
+    fg = themes.map((t, idx) => (idx === 0 && defaultColor ? '' : `${cssVariablePrefix + t.color}:`) + themeRegs[idx].fg).join(';')
+    bg = themes.map((t, idx) => (idx === 0 && defaultColor ? '' : `${cssVariablePrefix + t.color}-bg:`) + themeRegs[idx].bg).join(';')
     themeName = `shiki-themes ${themeRegs.map(t => t.name).join(' ')}`
     rootStyle = defaultColor ? undefined : [fg, bg].join(';')
   }
@@ -108,6 +82,48 @@ export function codeToHast(
   })
 }
 
+/**
+ *
+ */
+function flattenToken(
+  merged: ThemedTokenWithVariants,
+  variantsOrder: string[],
+  cssVariablePrefix: string,
+  defaultColor: string | boolean,
+) {
+  const token: ThemedToken = {
+    content: merged.content,
+    explanation: merged.explanation,
+  }
+
+  const styles = variantsOrder.map(t => getTokenStyleObject(merged.variants[t]))
+
+  // Get all style keys, for themes that missing some style, we put `inherit` to override as needed
+  const styleKeys = new Set(styles.flatMap(t => Object.keys(t)))
+  const mergedStyles = styles.reduce((acc, cur, idx) => {
+    for (const key of styleKeys) {
+      const value = cur[key] || 'inherit'
+
+      if (idx === 0 && defaultColor) {
+        acc[key] = value
+      }
+      else {
+        const varKey = cssVariablePrefix + variantsOrder[idx] + (key === 'color' ? '' : `-${key}`)
+        if (acc[key])
+          acc[key] += `;${varKey}:${value}`
+        else
+          acc[key] = `${varKey}:${value}`
+      }
+    }
+    return acc
+  }, {} as Record<string, string>)
+
+  token.htmlStyle = defaultColor
+    ? stringifyTokenStyle(mergedStyles)
+    : Object.values(mergedStyles).join(';')
+  return token
+}
+
 export function tokensToHast(
   tokens: ThemedToken[][],
   options: HtmlRendererOptions,
@@ -190,7 +206,7 @@ export function tokensToHast(
         children: [{ type: 'text', value: token.content }],
       }
 
-      const style = token.htmlStyle || stringifyTokenStyle(getTokenStyles(token))
+      const style = token.htmlStyle || stringifyTokenStyle(getTokenStyleObject(token))
       if (style)
         tokenNode.properties.style = style
 
@@ -223,7 +239,7 @@ export function tokensToHast(
   return result
 }
 
-function getTokenStyles(token: ThemedToken) {
+function getTokenStyleObject(token: TokenStyles) {
   const styles: Record<string, string> = {}
   if (token.color)
     styles.color = token.color
diff --git a/packages/shikiji/src/core/renderer-html-themes.ts b/packages/shikiji/src/core/renderer-html-themes.ts
@@ -1,30 +1,54 @@
-import type { CodeToTokensWithThemesOptions, ShikiContext, ThemedToken } from '../types'
+import type {
+  CodeToTokensWithThemesOptions,
+  ShikiContext,
+  ThemedToken,
+  ThemedTokenWithVariants,
+} from '../types'
 import { codeToThemedTokens } from './tokenizer'
 
 /**
- * Get tokens with multiple themes, with synced
+ * Get tokens with multiple themes
  */
 export function codeToTokensWithThemes(
   context: ShikiContext,
   code: string,
   options: CodeToTokensWithThemesOptions,
 ) {
   const themes = Object.entries(options.themes)
-    .filter(i => i[1]) as [string, string][]
+    .filter(i => i[1])
+    .map(i => ({ color: i[0], theme: i[1]! }))
 
   const tokens = syncThemesTokenization(
     ...themes.map(t => codeToThemedTokens(context, code, {
       ...options,
-      theme: t[1],
+      theme: t.theme,
       includeExplanation: false,
     })),
   )
 
-  return themes.map(([color, theme], idx) => [
-    color,
-    theme,
-    tokens[idx],
-  ] as [string, string, ThemedToken[][]])
+  const mergedTokens: ThemedTokenWithVariants[][] = tokens[0]
+    .map((line, lineIdx) => line
+      .map((_token, tokenIdx) => {
+        const mergedToken: ThemedTokenWithVariants = {
+          content: _token.content,
+          variants: {},
+        }
+
+        tokens.forEach((t, themeIdx) => {
+          const {
+            content: _,
+            explanation: __,
+            ...styles
+          } = t[lineIdx][tokenIdx]
+
+          mergedToken.variants[themes[themeIdx].color] = styles
+        })
+
+        return mergedToken
+      }),
+    )
+
+  return mergedTokens
 }
 
 /**
diff --git a/packages/shikiji/src/types.ts b/packages/shikiji/src/types.ts
@@ -54,40 +54,84 @@ export interface ShikiContext {
 }
 
 export interface HighlighterGeneric<BundledLangKeys extends string, BundledThemeKeys extends string> {
+  /**
+   * Get highlighted code in HTML string
+   */
   codeToHtml(
     code: string,
     options: CodeToHastOptions<ResolveBundleKey<BundledLangKeys>, ResolveBundleKey<BundledThemeKeys>>
   ): string
+  /**
+   * Get highlighted code in HAST.
+   * @see https://github.com/syntax-tree/hast
+   */
+  codeToHast(
+    code: string,
+    options: CodeToHastOptions<ResolveBundleKey<BundledLangKeys>, ResolveBundleKey<BundledThemeKeys>>
+  ): Root
+  /**
+   * Get highlighted code in tokens.
+   * @returns A 2D array of tokens, first dimension is lines, second dimension is tokens in a line.
+   */
   codeToThemedTokens(
     code: string,
     options: CodeToThemedTokensOptions<ResolveBundleKey<BundledLangKeys>, ResolveBundleKey<BundledThemeKeys>>
   ): ThemedToken[][]
+  /**
+   * Get highlighted code in tokens with multiple themes.
+   *
+   * Different from `codeToThemedTokens`, each token will have a `variants` property consisting of an object of color name to token styles.
+   *
+   * @returns A 2D array of tokens, first dimension is lines, second dimension is tokens in a line.
+   */
   codeToTokensWithThemes(
     code: string,
     options: CodeToTokensWithThemesOptions<ResolveBundleKey<BundledLangKeys>, ResolveBundleKey<BundledThemeKeys>>
-  ): [color: string, theme: string, tokens: ThemedToken[][]][]
-  codeToHast(
-    code: string,
-    options: CodeToHastOptions<ResolveBundleKey<BundledLangKeys>, ResolveBundleKey<BundledThemeKeys>>
-  ): Root
+  ): ThemedTokenWithVariants[][]
 
+  /**
+   * Load a theme to the highlighter, so later it can be used synchronously.
+   */
   loadTheme(...themes: (ThemeInput | BundledThemeKeys)[]): Promise<void>
+  /**
+   * Load a language to the highlighter, so later it can be used synchronously.
+   */
   loadLanguage(...langs: (LanguageInput | BundledLangKeys | SpecialLanguage)[]): Promise<void>
 
+  /**
+   * Get the theme registration object
+   */
   getTheme(name: string | ThemeRegistration | ThemeRegistrationRaw): ThemeRegistration
 
+  /**
+   * Get the names of loaded languages
+   *
+   * Special-handled languages like `text`, `plain` and `ansi` are not included.
+   */
   getLoadedLanguages(): string[]
+  /**
+   * Get the names of loaded themes
+   */
   getLoadedThemes(): string[]
 }
 
 export interface HighlighterCoreOptions {
+  /**
+   * Theme names, or theme registration objects to be loaded upfront.
+   */
   themes?: ThemeInput[]
+  /**
+   * Language names, or language registration objects to be loaded upfront.
+   */
   langs?: LanguageInput[]
   /**
    * Alias of languages
    * @example { 'my-lang': 'javascript' }
    */
   langAlias?: Record<string, string>
+  /**
+   * Load wasm file from a custom path or using a custom function.
+   */
   loadWasm?: OnigurumaLoadOptions | (() => Promise<OnigurumaLoadOptions>)
 }
 
@@ -426,31 +470,43 @@ export interface ThemedTokenExplanation {
  * }
  *
  */
-export interface ThemedToken {
+export interface ThemedToken extends TokenStyles, TokenBase {}
+
+export interface TokenBase {
   /**
    * The content of the token
    */
   content: string
   /**
-   * 6 or 8 digit hex code representation of the token's color
+   * Explanation of
+   *
+   * - token text's matching scopes
+   * - reason that token text is given a color (one matching scope matches a rule (scope -> color) in the theme)
    */
-  color?: string
+  explanation?: ThemedTokenExplanation[]
+}
+
+export interface TokenStyles {
   /**
-   * Override with custom inline style for HTML renderer.
-   * When specified, `color` will be ignored.
+   * 6 or 8 digit hex code representation of the token's color
    */
-  htmlStyle?: string
+  color?: string
   /**
    * Font style of token. Can be None/Italic/Bold/Underline
    */
   fontStyle?: FontStyle
   /**
-   * Explanation of
-   *
-   * - token text's matching scopes
-   * - reason that token text is given a color (one matching scope matches a rule (scope -> color) in the theme)
+   * Override with custom inline style for HTML renderer.
+   * When specified, `color` and `fontStyle` will be ignored.
    */
-  explanation?: ThemedTokenExplanation[]
+  htmlStyle?: string
+}
+
+export interface ThemedTokenWithVariants extends TokenBase {
+  /**
+   * An object of color name to token styles
+   */
+  variants: Record<string, TokenStyles>
 }
 
 export {}
diff --git a/packages/shikiji/test/themes.test.ts b/packages/shikiji/test/themes.test.ts
