feat(snapshot): add document symbols and folding support (#764)

Author: Gehbt

packages/extension/src/extension.ts

@@ -22,6 +22,9 @@ import { TestTree } from './testTree'
 import { getTestData, TestFile } from './testTreeData'
 import { clearCachedRuntime, debounce, showVitestError } from './utils'
 import './polyfills'
+import { SnapshotEntryTool } from './snapshot/tools'
+import { SnapshotDocumentSymbolProvider } from './snapshot/documentSymbolProvider'
+import { SnapshotFoldingRangeProvider } from './snapshot/foldingRangeProvider'
 
 export async function activate(context: vscode.ExtensionContext) {
   const extension = new VitestExtension(context)
@@ -320,6 +323,7 @@ class VitestExtension {
       'vitest.runtime',
       'deno.enabled',
     ]
+    const snapshotEntryTool = new SnapshotEntryTool()
 
     this.disposables = [
       vscode.workspace.onDidChangeConfiguration((event) => {
@@ -336,7 +340,7 @@ class VitestExtension {
         }),
       ),
       vscode.commands.registerCommand('vitest.openOutput', () => {
-        log.openOuput()
+        log.openOutput()
       }),
       vscode.commands.registerCommand('vitest.runRelatedTests', async (uri?: vscode.Uri) => {
         const currentUri = uri || vscode.window.activeTextEditor?.document.uri
@@ -537,6 +541,14 @@ class VitestExtension {
 
         await this.defineTestProfiles(false)
       }),
+      vscode.languages.registerDocumentSymbolProvider(
+        { language: 'vitest-snapshot' },
+        new SnapshotDocumentSymbolProvider(snapshotEntryTool),
+      ),
+      vscode.languages.registerFoldingRangeProvider(
+        { language: 'vitest-snapshot' },
+        new SnapshotFoldingRangeProvider(snapshotEntryTool),
+      ),
     ]
 
     // if the config changes, re-define all test profiles

packages/extension/src/log.ts

@@ -88,17 +88,17 @@ export const log = {
   workspaceError: (folder: string, ...args: any[]) => {
     log.error(`[Workspace ${folder}]`, ...args)
   },
-  openOuput() {
+  openOutput() {
     channel.show()
   },
 } as const
 
-let exitsts = false
+let exists = false
 function appendFile(log: string) {
-  if (!exitsts) {
+  if (!exists) {
     mkdirSync(dirname(logFile), { recursive: true })
     writeFileSync(logFile, '')
-    exitsts = true
+    exists = true
   }
   appendFileSync(logFile, `${log}\n`)
 }

packages/extension/src/snapshot/documentSymbolProvider.ts

@@ -0,0 +1,45 @@
+import * as vscode from 'vscode'
+import { createSnapshotSymbol, pushToDocumentSymbol, type SnapshotEntryTool } from './tools'
+
+export class SnapshotDocumentSymbolProvider implements vscode.DocumentSymbolProvider {
+  private latestUri: string | undefined = undefined
+  private latestVersion: number | undefined = undefined
+  latestDocumentSymbols: vscode.DocumentSymbol[] = []
+  constructor(private snapshotEntryTool: SnapshotEntryTool) {}
+  provideDocumentSymbols(
+    document: vscode.TextDocument,
+    token: vscode.CancellationToken,
+  ): vscode.ProviderResult<vscode.DocumentSymbol[]> {
+    if (this.latestUri === document.uri.toString() && this.latestVersion === document.version) {
+      return this.latestDocumentSymbols
+    }
+    this.snapshotEntryTool.process(document, document.uri.toString(), document.version, token)
+    if (token.isCancellationRequested) return null // cancelled
+
+    this.latestUri = document.uri.toString()
+    this.latestVersion = document.version
+
+    const documentSymbols: vscode.DocumentSymbol[] = []
+    forExportsSymbol: for (const entry of this.snapshotEntryTool.snapshotEntries) {
+      let currentLevel: vscode.DocumentSymbol[] = documentSymbols
+      let parent: vscode.DocumentSymbol[] | undefined
+
+      for (let i = 0; i < entry.breadcrumb.length; i++) {
+        const existingSymbol = currentLevel.at(-1)
+        if (!existingSymbol || existingSymbol.name !== entry.breadcrumb[i]) {
+          const newSymbol = createSnapshotSymbol(entry.breadcrumb[i], entry, i)
+          currentLevel.push(newSymbol)
+          i + 1 < entry.breadcrumb.length && pushToDocumentSymbol(newSymbol, entry, i + 1)
+          continue forExportsSymbol
+        }
+        parent = currentLevel
+        currentLevel = existingSymbol.children
+      }
+      // last level - all breadcrumbs matched, create duplicate leaf
+      ;(parent || documentSymbols).push(
+        createSnapshotSymbol(entry.breadcrumb.at(-1)!, entry, entry.breadcrumb.length - 1),
+      )
+    }
+    return (this.latestDocumentSymbols = documentSymbols)
+  }
+}

packages/extension/src/snapshot/foldingRangeProvider.ts

@@ -0,0 +1,34 @@
+import * as vscode from 'vscode'
+import { type SnapshotEntryTool } from './tools'
+
+export class SnapshotFoldingRangeProvider implements vscode.FoldingRangeProvider {
+  private latestUri: string | undefined = undefined
+  private latestVersion: number | undefined = undefined
+  latestFoldingRanges: vscode.FoldingRange[] = []
+  constructor(private snapshotEntryTool: SnapshotEntryTool) {}
+  provideFoldingRanges(
+    document: vscode.TextDocument,
+    _: vscode.FoldingContext,
+    token: vscode.CancellationToken,
+  ): vscode.ProviderResult<vscode.FoldingRange[]> {
+    if (this.latestUri === document.uri.toString() && this.latestVersion === document.version) {
+      return this.latestFoldingRanges
+    }
+    this.snapshotEntryTool.process(document, document.uri.toString(), document.version, token)
+    if (token.isCancellationRequested) return null // cancelled
+
+    this.latestUri = document.uri.toString()
+    this.latestVersion = document.version
+    const foldingRanges: vscode.FoldingRange[] = []
+    for (const symbol of this.snapshotEntryTool.snapshotEntries) {
+      foldingRanges.push(
+        new vscode.FoldingRange(
+          document.positionAt(symbol.start).line,
+          document.positionAt(symbol.end).line,
+          vscode.FoldingRangeKind.Region,
+        ),
+      )
+    }
+    return (this.latestFoldingRanges = foldingRanges)
+  }
+}

packages/extension/src/snapshot/tools.ts

@@ -0,0 +1,101 @@
+import * as vscode from 'vscode'
+
+const ExportSymbolRegex = /^exports\[`([^`]*)`\]/gm
+const RangeEndRegex = /`;$/m
+
+export interface SnapshotEntry {
+  name: string
+  breadcrumb: [...describeName: string[], itName: string]
+  start: number
+  end: number
+  fullRange: vscode.Range
+  keyRange: vscode.Range
+}
+
+export class SnapshotEntryTool {
+  private latestUri: string | undefined = undefined
+  private latestVersion: number | undefined = undefined
+  snapshotEntries: SnapshotEntry[] = []
+  process(
+    document: vscode.TextDocument,
+    uri: string,
+    version: number,
+    token: vscode.CancellationToken,
+  ): void {
+    let changeUri = false
+    let changeVersion = false
+    if (this.latestUri !== uri) {
+      this.latestUri = uri
+      this.latestVersion = version
+      changeUri = true
+      changeVersion = true
+    } else if (this.latestVersion !== version) {
+      this.latestVersion = version
+      changeVersion = true
+    }
+
+    if (!changeUri && !changeVersion) {
+      return // cached
+    } else {
+      // reset snapshotEntries
+      this.snapshotEntries = []
+    }
+    if (token.isCancellationRequested) return // cancelled
+    const text = document.getText()
+    const exportsSymbols = text.matchAll(ExportSymbolRegex) || []
+
+    for (const match of exportsSymbols) {
+      const name = match[1]
+      const snapshotDataStart = match.index
+      const snapshotDataEnd =
+        snapshotDataStart +
+        // find the nearest closing delimiter
+        (text.slice(snapshotDataStart).match(RangeEndRegex)?.index ??
+          // fallback to empty snapshot
+          'exports[`'.length + name.length + '`]'.length + ' = `'.length + '""'.length) +
+        '`;'.length
+
+      this.snapshotEntries.push({
+        name: name,
+        breadcrumb: name.split(' > ') as [...describeName: string[], itName: string],
+        start: snapshotDataStart,
+        end: snapshotDataEnd,
+        fullRange: new vscode.Range(
+          document.positionAt(snapshotDataStart),
+          document.positionAt(snapshotDataEnd),
+        ),
+        keyRange: new vscode.Range(
+          document.positionAt(snapshotDataStart + 'exports[`'.length),
+          document.positionAt(snapshotDataStart + 'exports[`'.length + name.length),
+        ),
+      })
+    }
+  }
+}
+
+export function createSnapshotSymbol(
+  name: string,
+  entry: SnapshotEntry,
+  index: number,
+): vscode.DocumentSymbol {
+  const isLastRound = index === entry.breadcrumb.length - 1
+  return new vscode.DocumentSymbol(
+    name,
+    isLastRound ? 'it' : 'describe',
+    vscode.SymbolKind.Function,
+    entry.fullRange,
+    entry.keyRange,
+  )
+}
+
+export function pushToDocumentSymbol(
+  parentDocumentSymbol: vscode.DocumentSymbol,
+  entry: SnapshotEntry,
+  startIndex: number = 1,
+): void {
+  for (let i = startIndex; i < entry.breadcrumb.length; i++) {
+    const newDocumentSymbol = createSnapshotSymbol(entry.breadcrumb[i], entry, i)
+    parentDocumentSymbol.children.push(newDocumentSymbol)
+    parentDocumentSymbol = newDocumentSymbol
+  }
+}

samples/basic/test/snapshot-case.test.ts

@@ -0,0 +1,52 @@
+import { describe, expect, it } from 'vitest'
+
+describe('fixture', () => {
+  describe('__fixtures__/file.spec.ts 1', () => {
+    it('snapshot', () => {
+      expect('').toMatchSnapshot()
+    })
+    it('snapshot_1', () => {
+      expect('').toMatchSnapshot()
+    })
+  })
+
+  describe('__fixtures__/file.spec.ts 2', () => {
+    it('snapshot_1', () => {
+      expect('').toMatchSnapshot()
+    })
+    it('snapshot_2', () => {
+      expect('').toMatchSnapshot()
+    })
+    it('snapshot', () => {
+      expect('\nsome content\n').toMatchSnapshot()
+    })
+  })
+
+  // same name it
+  describe('__fixtures__/file.spec.ts 2', () => {
+    it('snapshot_1', () => {
+      expect('').toMatchSnapshot()
+    })
+    it('snapshot_2', () => {
+      expect('').toMatchSnapshot()
+    })
+    it('snapshot', () => {
+      expect('').toMatchSnapshot()
+    })
+  })
+  // same name expect
+  it('snapshot_2', () => {
+    expect('').toMatchSnapshot()
+  })
+})
+
+describe('fixture2', () => {
+  it('__fixtures__/file.spec.ts 4', () => {
+    expect('').toMatchSnapshot()
+    expect('').toMatchSnapshot()
+  })
+})
+
+it('fixture2', () => {
+  expect('').toMatchSnapshot()
+})

test/e2e/utils/downloadSetup.ts

@@ -1,7 +1,7 @@
 import { download } from '@vscode/test-electron'
-import type { GlobalSetupContext } from 'vitest/node'
+import type { TestProject } from 'vitest/node'
 
-export default async function downloadVscode({ provide }: GlobalSetupContext) {
+export default async function downloadVscode({ provide }: TestProject) {
   if (process.env.VSCODE_E2E_DOWNLOAD_PATH)
     provide('executablePath', process.env.VSCODE_E2E_DOWNLOAD_PATH)
   else provide('executablePath', await download())