feat: introduce "show transformed module" command (#690)

Author: sheremet-va

package.json

@@ -80,13 +80,24 @@
         "title": "Reveal in Test Explorer",
         "command": "vitest.revealInTestExplorer",
         "category": "Vitest"
+      },
+      {
+        "title": "Open Transformed Module",
+        "command": "vitest.openTransformedModule",
+        "category": "Vitest"
       }
     ],
     "menus": {
       "editor/title/context": [
         {
           "command": "vitest.revealInTestExplorer",
-          "when": "vitest.testFiles && resourcePath in vitest.testFiles"
+          "when": "vitest.testFiles && resourcePath in vitest.testFiles",
+          "group": "vitest"
+        },
+        {
+          "command": "vitest.openTransformedModule",
+          "when": "vitest.environmentsSupported",
+          "group": "vitest"
         }
       ],
       "commandPalette": [

packages/extension/src/api.ts

@@ -43,6 +43,17 @@ export class VitestAPI {
     return this.api.forEach(callback)
   }
 
+  getModuleEnvironments(moduleId: string) {
+    return Promise.all(
+      this.api.map(async (api) => {
+        return {
+          api,
+          projects: await api.getModuleEnvironments(moduleId),
+        }
+      }),
+    )
+  }
+
   get folderAPIs() {
     return this.api
   }
@@ -111,6 +122,14 @@ export class VitestFolderAPI {
     return this.pkg
   }
 
+  getTransformedModule(project: string, environment: string, moduleId: string) {
+    return this.meta.rpc.getTransformedModule(project, environment, moduleId)
+  }
+
+  async getModuleEnvironments(moduleId: string) {
+    return this.meta.rpc.getModuleEnvironments(moduleId)
+  }
+
   async runFiles(specs?: ExtensionTestSpecification[] | string[], testNamePatern?: string) {
     await this.meta.rpc.runTests(normalizeSpecs(specs), testNamePatern)
   }

packages/extension/src/api/ws.ts

@@ -4,6 +4,7 @@ import type { ResolvedMeta } from '../api'
 import type { VitestPackage } from './pkg'
 import { pathToFileURL } from 'node:url'
 import { gte } from 'semver'
+import vscode from 'vscode'
 import { getConfig } from '../config'
 import { browserSetupFilePath, finalCoverageFileName, setupFilePath } from '../constants'
 import { log } from '../log'
@@ -72,6 +73,13 @@ export function onWsConnection(
         log.verbose?.('[API]', 'Vitest WebSocket connection closed, cannot call RPC anymore.')
         api.$close()
       })
+      if (!message.legacy) {
+        vscode.commands.executeCommand(
+          'setContext',
+          'vitest.environmentsSupported',
+          true,
+        )
+      }
       onStart({
         rpc: api,
         workspaceSource: message.workspaceSource,

packages/extension/src/extension.ts

@@ -12,6 +12,7 @@ import { DebugManager, debugTests } from './debug'
 import { ExtensionDiagnostic } from './diagnostic'
 import { log } from './log'
 import { TestRunner } from './runner'
+import { SchemaProvider } from './schemaProvider'
 import { TagsManager } from './tagsManager'
 import { TestTree } from './testTree'
 import { getTestData, TestFile } from './testTreeData'
@@ -41,6 +42,7 @@ class VitestExtension {
   private disposables: vscode.Disposable[] = []
   private diagnostic: ExtensionDiagnostic | undefined
   private debugManager: DebugManager
+  private schemaProvider: SchemaProvider
 
   /** @internal */
   _debugDisposable: vscode.Disposable | undefined
@@ -58,6 +60,12 @@ class VitestExtension {
     this.testTree = new TestTree(this.testController, this.loadingTestItem)
     this.tagsManager = new TagsManager(this.testTree)
     this.debugManager = new DebugManager()
+    this.schemaProvider = new SchemaProvider(
+      async (apiId, project, environment, file) => {
+        const api = this.api?.folderAPIs.find(a => a.id === apiId)
+        return api?.getTransformedModule(project, environment, file) ?? null
+      },
+    )
   }
 
   private _defineTestProfilePromise: Promise<void> | undefined
@@ -344,6 +352,51 @@ class VitestExtension {
         const tokenSource = new vscode.CancellationTokenSource()
         await profile.runHandler(request, tokenSource.token)
       }),
+      vscode.commands.registerCommand('vitest.openTransformedModule', async (uri: vscode.Uri | undefined) => {
+        const currentUri = uri || vscode.window.activeTextEditor?.document.uri
+        if (!this.api || !currentUri || currentUri.scheme === 'vitest-transform') {
+          return
+        }
+        const environments = await this.api.getModuleEnvironments(currentUri.fsPath)
+        const options = environments.map(({ api, projects }) => {
+          return projects.map((project) => {
+            return project.environments.map((environment) => {
+              let label = ''
+              if (environments.length > 1) {
+                label += `${api.prefix}: `
+              }
+              if (project.name) {
+                label += `[${project.name}] `
+              }
+              label += environment
+              return {
+                label,
+                uriParts: [api.id, project.name, environment],
+              }
+            })
+          })
+        }).flat(2)
+        if (options.length === 0) {
+          vscode.window.showWarningMessage('All module graphs are empty, nothing to show.')
+          return
+        }
+        const pick = options.length === 1 ? options[0] : await vscode.window.showQuickPick(options)
+        if (!pick) {
+          return
+        }
+        try {
+          const [apiId, projectName, environment] = pick.uriParts
+          const uri = vscode.Uri.parse(
+            `vitest-transform://${currentUri.fsPath}.js?apiId=${apiId}&project=${projectName}&environment=${environment}`,
+          )
+          const doc = await vscode.workspace.openTextDocument(uri)
+          await vscode.window.showTextDocument(doc, { preview: false })
+        }
+        catch (err) {
+          log.error(err)
+          vscode.window.showErrorMessage(`Vitest: The file was not processed by Vite yet. Try running the tests first${options.length > 1 ? ' or select a different environment' : ''}.`)
+        }
+      }),
     ]
 
     // if the config changes, re-define all test profiles
@@ -429,6 +482,7 @@ class VitestExtension {
     this.testTree.dispose()
     this.tagsManager.dispose()
     this.testController.dispose()
+    this.schemaProvider.dispose()
     this.runProfiles.forEach(profile => profile.dispose())
     this.runProfiles.clear()
     this.disposables.forEach(d => d.dispose())

packages/extension/src/schemaProvider.ts

@@ -0,0 +1,31 @@
+import * as vscode from 'vscode'
+
+export class SchemaProvider implements vscode.TextDocumentContentProvider, vscode.Disposable {
+  private disposables: vscode.Disposable[] = []
+
+  constructor(
+    private getTransformedModule: (apiId: string, project: string, environment: string, file: string) => Promise<string | null>,
+  ) {
+    this.disposables.push(vscode.workspace.registerTextDocumentContentProvider('vitest-transform', this))
+  }
+
+  async provideTextDocumentContent(uri: vscode.Uri): Promise<string> {
+    const searchParams = new URLSearchParams(uri.query)
+    const apiId = searchParams.get('apiId')
+    const project = searchParams.get('project')
+    const environment = searchParams.get('environment')
+    if (apiId == null || project == null || environment == null) {
+      throw new Error(`Cannot parse the schema: ${uri.toString()}`)
+    }
+    const fsPath = uri.fsPath.slice(0, -3) // remove .js
+    const content = await this.getTransformedModule(apiId, project, environment, fsPath)
+    if (content == null) {
+      throw new Error(`The file ${fsPath} was not processed by Vite yet.`)
+    }
+    return content
+  }
+
+  dispose() {
+    this.disposables.forEach(d => d.dispose())
+  }
+}

packages/extension/src/worker/index.ts

@@ -25,8 +25,8 @@ emitter.on('message', async function onMessage(message: any) {
         pathToFileURL(normalizeDriveLetter(data.meta.vitestNodePath)).toString()
       ) as typeof import('vitest/node')
 
-      const isOld = !vitestModule.version || (Number(vitestModule.version[0]) < 4)
-      const workerName = isOld ? './workerLegacy.js' : './workerNew.js'
+      const isLegacy = !vitestModule.version || (Number(vitestModule.version[0]) < 4)
+      const workerName = isLegacy ? './workerLegacy.js' : './workerNew.js'
       const workerPath = pathToFileURL(join(__dirname, workerName))
       const initModule = await import(workerPath.toString())
 
@@ -53,7 +53,7 @@ emitter.on('message', async function onMessage(message: any) {
       )
       worker.initRpc(rpc)
       reporter.initRpc(rpc)
-      emitter.ready(configs, workspaceSource, isOld)
+      emitter.ready(configs, workspaceSource, isLegacy)
     }
     catch (err: any) {
       emitter.error(err)

packages/shared/src/index.ts

@@ -37,6 +37,11 @@ export interface ExtensionUserConsoleLog extends UserConsoleLog {
   }
 }
 
+export interface ExtensionEnvironment {
+  name: string
+  environments: string[]
+}
+
 export interface ExtensionWorkerTransport {
   getFiles: () => Promise<ExtensionTestFileSpecification[]>
   collectTests: (testFile: ExtensionTestSpecification[]) => Promise<void>
@@ -58,6 +63,8 @@ export interface ExtensionWorkerTransport {
   onFilesChanged: (files: string[]) => void
 
   initRpc: (rpc: VitestWorkerRPC) => void
+  getModuleEnvironments: (moduleId: string) => ExtensionEnvironment[]
+  getTransformedModule: (project: string, environment: string, moduleId: string) => string | null
 
   onBrowserDebug: (fulfilled: boolean) => void
 }

packages/worker-legacy/src/worker.ts

@@ -425,4 +425,13 @@ export class ExtensionWorker implements ExtensionWorkerTransport {
   onBrowserDebug(fulfilled: boolean) {
     ExtensionWorker.emitter.emit('onBrowserDebug', fulfilled)
   }
+
+  // TODO:(?) -- if environments are supported
+  getModuleEnvironments() {
+    return []
+  }
+
+  getTransformedModule() {
+    return null
+  }
 }

packages/worker/src/worker.ts

@@ -1,4 +1,5 @@
 import type {
+  ExtensionEnvironment,
   ExtensionTestFileSpecification,
   ExtensionTestSpecification,
   ExtensionWorkerTransport,
@@ -103,6 +104,35 @@ export class ExtensionWorker implements ExtensionWorkerTransport {
     this.runner.initRpc(rpc)
   }
 
+  getModuleEnvironments(moduleId: string): ExtensionEnvironment[] {
+    return this.vitest.projects.map((project) => {
+      const environments = new Set<string>()
+      for (const name in project.vite.environments) {
+        const environment = project.vite.environments[name]
+        const nodes = [...environment.moduleGraph.getModulesByFile(moduleId) || []]
+        if (nodes.some(n => n.transformResult)) {
+          environments.add(name)
+        }
+      }
+      return {
+        name: project.name,
+        environments: Array.from(environments),
+      }
+    })
+  }
+
+  getTransformedModule(projectName: string, environmentName: string, moduleId: string) {
+    const project = this.vitest.projects.find(p => p.name === projectName)
+    const environment = environmentName === '__browser__'
+      ? project?.browser?.vite?.environments.client
+      : project?.vite.environments[environmentName]
+    const files = environment?.moduleGraph.getModulesByFile(moduleId)
+    if (!files || !files.size) {
+      return null
+    }
+    return files.values().next().value?.transformResult?.code ?? null
+  }
+
   onBrowserDebug(fulfilled: boolean) {
     ExtensionWorker.emitter.emit('onBrowserDebug', fulfilled)
   }