feat(twoslash): support @includes (#737)

Co-authored-by: Anthony Fu <github@antfu.me>

packages/twoslash/src/core.ts

@@ -9,6 +9,7 @@ import type { Element, ElementContent, Text } from 'hast'
 import { splitTokens } from '@shikijs/core'
 import type { TransformerTwoslashOptions, TwoslashRenderer, TwoslashShikiFunction, TwoslashShikiReturn } from './types'
 import { ShikiTwoslashError } from './error'
+import { TwoslashIncludesManager, parseIncludeMeta } from './includes'
 
 export * from './types'
 export * from './renderer-rich'
@@ -39,6 +40,7 @@ export function createTransformerFactory(
       explicitTrigger = false,
       renderer = defaultRenderer,
       throws = true,
+      includesMap = new Map(),
     } = options
 
     const onTwoslashError = options.onTwoslashError || (
@@ -66,6 +68,9 @@ export function createTransformerFactory(
     const map = new WeakMap<ShikiTransformerContextMeta, TwoslashShikiReturn>()
 
     const filter = options.filter || ((lang, _, options) => langs.includes(lang) && (!explicitTrigger || trigger.test(options.meta?.__raw || '')))
+
+    const includes = new TwoslashIncludesManager(includesMap)
+
     return {
       preprocess(code) {
         let lang = this.options.lang
@@ -74,7 +79,14 @@ export function createTransformerFactory(
 
         if (filter(lang, code, this.options)) {
           try {
-            const twoslash = (twoslasher as TwoslashShikiFunction)(code, lang, twoslashOptions)
+            const include = parseIncludeMeta(this.options.meta?.__raw)
+
+            if (include)
+              includes.add(include, code)
+
+            const codeWithIncludes = includes.applyInclude(code)
+
+            const twoslash = (twoslasher as TwoslashShikiFunction)(codeWithIncludes, lang, twoslashOptions)
             map.set(this.meta, twoslash)
             this.meta.twoslash = twoslash
             this.options.lang = twoslash.meta?.extension || lang

packages/twoslash/src/includes.ts

@@ -0,0 +1,73 @@
+export class TwoslashIncludesManager {
+  constructor(
+    public map: Map<string, string> = new Map(),
+  ) {}
+
+  add(name: string, code: string) {
+    const lines: string[] = []
+
+    code.split('\n').forEach((l, _i) => {
+      const trimmed = l.trim()
+
+      if (trimmed.startsWith('// - ')) {
+        const key = trimmed.split('// - ')[1].split(' ')[0]
+        this.map.set(`${name}-${key}`, lines.join('\n'))
+      }
+      else {
+        lines.push(l)
+      }
+    })
+    this.map.set(name, lines.join('\n'))
+  }
+
+  applyInclude(code: string) {
+    const reMarker = /\/\/ @include: (.*)$/gm
+
+    // Basically run a regex over the code replacing any // @include: thing with
+    // 'thing' from the map
+
+    // const toReplace: [index:number, length: number, str: string][] = []
+    const toReplace: [number, number, string][] = []
+
+    for (const match of code.matchAll(reMarker)) {
+      const key = match[1]
+      const replaceWith = this.map.get(key)
+
+      if (!replaceWith) {
+        const msg = `Could not find an include with the key: '${key}'.\nThere is: ${Array.from(this.map.keys())}.`
+        throw new Error(msg)
+      }
+      else {
+        toReplace.push([match.index, match[0].length, replaceWith])
+      }
+    }
+
+    let newCode = code.toString()
+    // Go backwards through the found changes so that we can retain index position
+    toReplace
+      .reverse()
+      .forEach((r) => {
+        newCode = newCode.slice(0, r[0]) + r[2] + newCode.slice(r[0] + r[1])
+      })
+    return newCode
+  }
+}
+
+/**
+ * An "include [name]" segment in a raw meta string is a sequence of words,
+ * possibly connected by dashes, following "include " and ending at a word boundary.
+ */
+const INCLUDE_META_REGEX = /include\s+([\w-]+)\b.*/
+
+/**
+ * Given a raw meta string for code block like 'twoslash include main-hello-world meta=miscellaneous',
+ * capture the name of the reusable code block as "main-hello-world", and ignore anything
+ * before and after this segment.
+ */
+export function parseIncludeMeta(meta?: string): string | null {
+  if (!meta)
+    return null
+
+  const match = meta.match(INCLUDE_META_REGEX)
+  return match?.[1] ?? null
+}

packages/twoslash/src/types.ts

@@ -48,6 +48,11 @@ export interface TransformerTwoslashOptions {
    * Custom renderers to decide how each info should be rendered
    */
   renderer?: TwoslashRenderer
+  /**
+   * A map to store code for `@include` directive
+   * Provide your own instance if you want to clear the map between each transformation
+   */
+  includesMap?: Map<string, string>
   /**
    * Strictly throw when there is an error
    * @default true

packages/twoslash/test/includes.test.ts

@@ -0,0 +1,117 @@
+import { expect, it } from 'vitest'
+import { codeToHtml } from 'shiki'
+import { TwoslashIncludesManager } from '../src/includes'
+import { rendererRich, transformerTwoslash } from '../src'
+
+const styleTag = `
+<link rel="stylesheet" href="../../../style-rich.css" />
+<style>
+.dark .shiki,
+.dark .shiki span {
+  color: var(--shiki-dark, inherit);
+  background-color: var(--shiki-dark-bg, inherit);
+  --twoslash-popup-bg: var(--shiki-dark-bg, inherit);
+}
+
+html:not(.dark) .shiki,
+html:not(.dark) .shiki span {
+  color: var(--shiki-light, inherit);
+  background-color: var(--shiki-light-bg, inherit);
+  --twoslash-popup-bg: var(--shiki-light-bg, inherit);
+}
+</style>
+`
+
+const multiExample = `
+const a = 1
+// - 1
+const b = 2
+// - 2
+const c = 3
+`
+
+it('creates a set of examples', () => {
+  const manager = new TwoslashIncludesManager()
+  manager.add('main', multiExample)
+  expect(manager.map.size === 3)
+
+  expect(manager.map.get('main')).toContain('const c')
+  expect(manager.map.get('main-1')).toContain('const a = 1')
+  expect(manager.map.get('main-2')).toContain('const b = 2')
+})
+
+it('replaces the code', () => {
+  const manager = new TwoslashIncludesManager()
+  manager.add('main', multiExample)
+  expect(manager.map.size === 3)
+
+  const sample = `// @include: main`
+  const replaced = manager.applyInclude(sample)
+  expect(replaced).toMatchInlineSnapshot(`
+    "
+    const a = 1
+    const b = 2
+    const c = 3
+    "
+  `)
+})
+
+it('throws an error if key not found', () => {
+  const manager = new TwoslashIncludesManager()
+
+  const sample = `// @include: main`
+  expect(() => manager.applyInclude(sample)).toThrow()
+})
+
+it('replaces @include directives with previously transformed code blocks', async () => {
+  const main = `
+export const hello = { str: "world" };
+`.trim()
+
+  /**
+   * The @noErrors directive allows the code above the ^| to be invalid,
+   * i.e. so it can demonstrate what a partial autocomplete looks like.
+   */
+  const code = `
+// @include: main
+// @noErrors
+
+hello.
+//    ^|
+`.trim()
+
+  /**
+   * Replacing @include directives only renders nicely rendererRich?
+   */
+  const transformer = transformerTwoslash({
+    renderer: rendererRich(),
+  })
+
+  const htmlMain = await codeToHtml(main, {
+    lang: 'ts',
+    themes: {
+      dark: 'vitesse-dark',
+      light: 'vitesse-light',
+    },
+    defaultColor: false,
+    transformers: [transformer],
+    meta: {
+      __raw: 'include main',
+    },
+  })
+
+  expect(styleTag + htmlMain).toMatchFileSnapshot('./out/includes/main.html')
+
+  const html = await codeToHtml(code, {
+    lang: 'ts',
+    themes: {
+      dark: 'vitesse-dark',
+      light: 'vitesse-light',
+    },
+    transformers: [transformer],
+  })
+
+  expect(styleTag + html).toMatchFileSnapshot(
+    './out/includes/replaced_directives.html',
+  )
+})