feat(shortcuts): guess embedded language and auto load in shortcuts (#932)

Author: antfu

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

@@ -134,6 +134,12 @@ export function createdBundledHighlighter<BundledLangs extends string, BundledTh
       loadTheme(...themes) {
         return core.loadTheme(...themes.map(resolveTheme))
       },
+      getBundledLanguages() {
+        return bundledLanguages
+      },
+      getBundledThemes() {
+        return bundledThemes
+      },
     }
   }
 
@@ -224,53 +230,58 @@ export function makeSingletonHighlighter<L extends string, T extends string>(
   return getSingletonHighlighter
 }
 
+export interface CreateSingletonShorthandsOptions<L extends string, T extends string> {
+  /**
+   * A custom function to guess embedded languages to be loaded.
+   */
+  guessEmbeddedLanguages?: (code: string, lang: string | undefined, highlighter: HighlighterGeneric<L, T>) => Awaitable<string[] | undefined>
+}
+
 export function createSingletonShorthands<L extends string, T extends string>(
   createHighlighter: CreateHighlighterFactory<L, T>,
+  config?: CreateSingletonShorthandsOptions<L, T>,
 ): ShorthandsBundle<L, T> {
   const getSingletonHighlighter = makeSingletonHighlighter(createHighlighter)
 
+  async function get(code: string, options: CodeToTokensOptions<L, T> | CodeToHastOptions<L, T>): Promise<HighlighterGeneric<L, T>> {
+    const shiki = await getSingletonHighlighter({
+      langs: [options.lang as L],
+      themes: ('theme' in options ? [options.theme] : Object.values(options.themes)) as T[],
+    })
+    const langs = await config?.guessEmbeddedLanguages?.(code, options.lang, shiki) as L[]
+    if (langs) {
+      await shiki.loadLanguage(...langs)
+    }
+    return shiki
+  }
+
   return {
     getSingletonHighlighter(options) {
       return getSingletonHighlighter(options)
     },
 
     async codeToHtml(code, options) {
-      const shiki = await getSingletonHighlighter({
-        langs: [options.lang as L],
-        themes: ('theme' in options ? [options.theme] : Object.values(options.themes)) as T[],
-      })
+      const shiki = await get(code, options)
       return shiki.codeToHtml(code, options)
     },
 
     async codeToHast(code, options) {
-      const shiki = await getSingletonHighlighter({
-        langs: [options.lang as L],
-        themes: ('theme' in options ? [options.theme] : Object.values(options.themes)) as T[],
-      })
+      const shiki = await get(code, options)
       return shiki.codeToHast(code, options)
     },
 
     async codeToTokens(code, options) {
-      const shiki = await getSingletonHighlighter({
-        langs: [options.lang as L],
-        themes: ('theme' in options ? [options.theme] : Object.values(options.themes)) as T[],
-      })
+      const shiki = await get(code, options)
       return shiki.codeToTokens(code, options)
     },
 
     async codeToTokensBase(code, options) {
-      const shiki = await getSingletonHighlighter({
-        langs: [options.lang as L],
-        themes: [options.theme as T],
-      })
+      const shiki = await get(code, options)
       return shiki.codeToTokensBase(code, options)
     },
 
     async codeToTokensWithThemes(code, options) {
-      const shiki = await getSingletonHighlighter({
-        langs: [options.lang as L],
-        themes: Object.values(options.themes).filter(Boolean) as T[],
-      })
+      const shiki = await get(code, options)
       return shiki.codeToTokensWithThemes(code, options)
     },
 

packages/core/src/constructors/highlighter.ts

@@ -26,6 +26,8 @@ export async function createHighlighterCore(options: HighlighterCoreOptions<fals
     codeToTokens: (code, options) => codeToTokens(internal, code, options),
     codeToHast: (code, options) => codeToHast(internal, code, options),
     codeToHtml: (code, options) => codeToHtml(internal, code, options),
+    getBundledLanguages: () => ({}),
+    getBundledThemes: () => ({}),
     ...internal,
     getInternalContext: () => internal,
   }
@@ -49,6 +51,8 @@ export function createHighlighterCoreSync(options: HighlighterCoreOptions<true>)
     codeToTokens: (code, options) => codeToTokens(internal, code, options),
     codeToHast: (code, options) => codeToHast(internal, code, options),
     codeToHtml: (code, options) => codeToHtml(internal, code, options),
+    getBundledLanguages: () => ({}),
+    getBundledThemes: () => ({}),
     ...internal,
     getInternalContext: () => internal,
   }

packages/core/src/highlight/code-to-tokens-base.ts

@@ -185,7 +185,7 @@ function _tokenizeWithTheme(
     let tokensWithScopesIndex
 
     if (options.includeExplanation) {
-      resultWithScopes = grammar.tokenizeLine(line, stateStack)
+      resultWithScopes = grammar.tokenizeLine(line, stateStack, tokenizeTimeLimit)
       tokensWithScopes = resultWithScopes.tokens
       tokensWithScopesIndex = 0
     }

packages/core/src/textmate/registry.ts

@@ -67,7 +67,6 @@ export class Registry extends TextMateRegistry {
       this._textmateThemeCache.set(theme, textmateTheme)
     }
 
-    // @ts-expect-error Access private `_syncRegistry`, but should work in runtime
     this._syncRegistry.setTheme(textmateTheme)
   }
 
@@ -100,7 +99,6 @@ export class Registry extends TextMateRegistry {
       unbalancedBracketSelectors: lang.unbalancedBracketSelectors || [],
     }
 
-    // @ts-expect-error Private members, set this to override the previous grammar (that can be a stub)
     this._syncRegistry._rawGrammars.set(lang.scopeName, lang)
     const g = this.loadGrammarWithConfiguration(lang.scopeName, 1, grammarConfig) as Grammar
     g.name = lang.name
@@ -119,9 +117,7 @@ export class Registry extends TextMateRegistry {
         this._resolvedGrammars.delete(e.name)
         // Reset cache
         this._loadedLanguagesCache = null
-        // @ts-expect-error clear cache
         this._syncRegistry?._injectionGrammars?.delete(e.scopeName)
-        // @ts-expect-error clear cache
         this._syncRegistry?._grammars?.delete(e.scopeName)
         this.loadLanguage(this._langMap.get(e.name)!)
       }

packages/shiki/src/bundle-full.ts

@@ -4,6 +4,7 @@ import type { BundledLanguage } from './langs-bundle-full'
 import type { BundledTheme } from './themes'
 import { createdBundledHighlighter, createSingletonShorthands, warnDeprecated } from './core'
 import { createOnigurumaEngine } from './engine-oniguruma'
+import { guessEmbeddedLanguages } from './guess'
 import { bundledLanguages } from './langs-bundle-full'
 import { bundledThemes } from './themes'
 
@@ -46,6 +47,7 @@ export const {
   BundledTheme
 >(
   createHighlighter,
+  { guessEmbeddedLanguages },
 )
 
 /**

packages/shiki/src/bundle-web.ts

@@ -4,6 +4,7 @@ import type { BundledLanguage } from './langs-bundle-web'
 import type { BundledTheme } from './themes'
 import { createdBundledHighlighter, createSingletonShorthands, warnDeprecated } from './core'
 import { createOnigurumaEngine } from './engine-oniguruma'
+import { guessEmbeddedLanguages } from './guess'
 import { bundledLanguages } from './langs-bundle-web'
 import { bundledThemes } from './themes'
 
@@ -46,6 +47,7 @@ export const {
   BundledTheme
 >(
   createHighlighter,
+  { guessEmbeddedLanguages },
 )
 
 /**

packages/shiki/src/guess.ts

@@ -0,0 +1,26 @@
+import type { HighlighterGeneric } from '@shikijs/types'
+
+export function guessEmbeddedLanguages(
+  code: string,
+  _lang: string | undefined,
+  shiki: HighlighterGeneric<any, any>,
+): string[] {
+  const langs = new Set<string>()
+  // For HTML code blocks like Vue SFC
+  for (const match of code.matchAll(/lang=["']([\w-]+)["']/g)) {
+    langs.add(match[1])
+  }
+  // For markdown code blocks
+  for (const match of code.matchAll(/(?:```|~~~)([\w-]+)/g)) {
+    langs.add(match[1])
+  }
+  // For latex
+  for (const match of code.matchAll(/\\begin\{([\w-]+)\}/g)) {
+    langs.add(match[1])
+  }
+
+  // Only include known languages
+  const bundle = shiki.getBundledLanguages()
+  return Array.from(langs)
+    .filter(l => l && bundle[l])
+}

packages/shiki/test/out/shorthand-markdown1.html

@@ -0,0 +1,76 @@
+import { expect, it } from 'vitest'
+import { codeToHtml, getSingletonHighlighter } from '../src'
+
+const inputMarkdown1 = `
+This is a markdown file
+
+\`\`\`js
+console.log("hello")
+\`\`\`
+
+~~~pug
+div
+  p hello
+~~~
+
+Even those grammars in markdown are lazy loaded, \`codeToHtml\` shorthand should capture them and load automatically.
+`
+
+const inputMarkdown2 = `
+Some other languages
+
+\`\`\`js
+console.log("hello")
+\`\`\`
+
+~~~python
+print("hello")
+~~~
+
+\`\`\`html
+<div class="foo">bar</div>
+<style>
+  .foo {
+    color: red;
+  }
+</style>
+\`\`\`
+`
+
+it('codeToHtml', async () => {
+  const highlighter = await getSingletonHighlighter()
+  expect(highlighter.getLoadedLanguages())
+    .toEqual([])
+
+  await expect(await codeToHtml(inputMarkdown1, { lang: 'markdown', theme: 'vitesse-light' }))
+    .toMatchFileSnapshot(`out/shorthand-markdown1.html`)
+
+  expect.soft(highlighter.getLoadedLanguages())
+    .toContain('javascript')
+  expect.soft(highlighter.getLoadedLanguages())
+    .toContain('pug')
+
+  await expect(await codeToHtml(inputMarkdown2, { lang: 'markdown', theme: 'vitesse-light' }))
+    .toMatchFileSnapshot(`out/shorthand-markdown2.html`)
+
+  expect.soft(highlighter.getLoadedLanguages())
+    .toContain('python')
+  expect.soft(highlighter.getLoadedLanguages())
+    .toContain('html')
+
+  expect.soft(highlighter.getLoadedLanguages())
+    .toMatchInlineSnapshot(`
+      [
+        "javascript",
+        "css",
+        "html",
+        "pug",
+        "python",
+        "markdown",
+        "md",
+        "js",
+        "jade",
+        "py",
+      ]
+    `)
+})

packages/shiki/test/out/shorthand-markdown2.html

@@ -143,6 +143,14 @@ export interface HighlighterGeneric<BundledLangKeys extends string, BundledTheme
    * @deprecated
    */
   getInternalContext: () => ShikiInternal
+  /**
+   * Get bundled languages object
+   */
+  getBundledLanguages: () => Record<BundledLangKeys, LanguageInput>
+  /**
+   * Get bundled themes object
+   */
+  getBundledThemes: () => Record<BundledThemeKeys, ThemeInput>
 }
 
 /**

packages/shiki/test/shorthands-markdown.test.ts

@@ -19,7 +19,7 @@ catalog:
   '@rollup/plugin-node-resolve': ^16.0.0
   '@rollup/plugin-replace': ^6.0.2
   '@rollup/plugin-terser': ^0.4.4
-  '@shikijs/vscode-textmate': ^10.0.1
+  '@shikijs/vscode-textmate': ^10.0.2
   '@types/fs-extra': ^11.0.4
   '@types/hast': ^3.0.4
   '@types/markdown-it': ^14.1.2