feat: add multi-env support for build (#148)

Author: sapphi-red

src/node/build.ts

@@ -1,3 +1,4 @@
+import type { Environment } from 'vite'
 import type { ModuleTransformInfo } from '../types'
 import type { InspectContext } from './context'
 import fs from 'node:fs/promises'
@@ -6,9 +7,36 @@ import process from 'node:process'
 import { hash } from 'ohash'
 import { DIR_CLIENT } from '../dirs'
 
-export async function generateBuild(
-  ctx: InspectContext,
+export interface EnvOrderHooks<Args extends readonly unknown[]> {
+  onFirst?: (...args: Args) => Promise<void>
+  onEach?: (...args: Args) => Promise<void>
+  onLast?: (...args: Args) => Promise<void>
+}
+export function createEnvOrderHooks<Args extends readonly unknown[]>(
+  environmentNames: string[],
+  { onFirst, onEach, onLast }: EnvOrderHooks<Args>,
 ) {
+  const remainingEnvs = new Set<string>(environmentNames)
+  let ranFirst = false
+  let ranLast = false
+
+  return async (envName: string, ...args: Args) => {
+    if (!ranFirst) {
+      ranFirst = true
+      await onFirst?.(...args)
+    }
+
+    remainingEnvs.delete(envName)
+    await onEach?.(...args)
+
+    if (!ranLast && remainingEnvs.size === 0) {
+      ranLast = true
+      await onLast?.(...args)
+    }
+  }
+}
+
+export function createBuildGenerator(ctx: InspectContext) {
   const {
     outputDir = '.vite-inspect',
   } = ctx.options
@@ -19,55 +47,62 @@ export async function generateBuild(
     : resolve(process.cwd(), outputDir)
   const reportsDir = join(targetDir, 'reports')
 
-  await fs.rm(targetDir, { recursive: true, force: true })
-  await fs.mkdir(reportsDir, { recursive: true })
-  await fs.cp(DIR_CLIENT, targetDir, { recursive: true })
-
-  await Promise.all([
-    fs.writeFile(
-      join(targetDir, 'index.html'),
-      (await fs.readFile(join(targetDir, 'index.html'), 'utf-8'))
-        .replace(
-          'data-vite-inspect-mode="DEV"',
-          'data-vite-inspect-mode="BUILD"',
+  return {
+    getOutputDir() {
+      return targetDir
+    },
+    async setupOutputDir() {
+      await fs.rm(targetDir, { recursive: true, force: true })
+      await fs.mkdir(reportsDir, { recursive: true })
+      await fs.cp(DIR_CLIENT, targetDir, { recursive: true })
+      await Promise.all([
+        fs.writeFile(
+          join(targetDir, 'index.html'),
+          (await fs.readFile(join(targetDir, 'index.html'), 'utf-8'))
+            .replace(
+              'data-vite-inspect-mode="DEV"',
+              'data-vite-inspect-mode="BUILD"',
+            ),
+        ),
+        writeJSON(
+          join(reportsDir, 'metadata.json'),
+          ctx.getMetadata(),
         ),
-    ),
-    writeJSON(
-      join(reportsDir, 'metadata.json'),
-      ctx.getMetadata(),
-    ),
-    ...[...ctx._idToInstances.values()]
-      .flatMap(v => [...v.environments.values()]
-        .map((e) => {
-          const key = `${v.id}-${e.env.name}`
+      ])
+    },
+    async generateForEnv(env: Environment) {
+      await Promise.all([...ctx._idToInstances.values()]
+        .filter(v => v.environments.has(env.name))
+        .map((v) => {
+          const e = v.environments.get(env.name)!
+          const key = `${v.id}-${env.name}`
           return [key, e] as const
+        })
+        .map(async ([key, env]) => {
+          await fs.mkdir(join(reportsDir, key, 'transforms'), { recursive: true })
+
+          return await Promise.all([
+            writeJSON(
+              join(reportsDir, key, 'modules.json'),
+              env.getModulesList(),
+            ),
+            writeJSON(
+              join(reportsDir, key, 'metric-plugins.json'),
+              env.getPluginMetrics(),
+            ),
+            ...Object.entries(env.data.transform)
+              .map(([id, info]) => writeJSON(
+                join(reportsDir, key, 'transforms', `${hash(id)}.json`),
+            <ModuleTransformInfo>{
+              resolvedId: id,
+              transforms: info,
+            },
+              )),
+          ])
         }),
       )
-      .map(async ([key, env]) => {
-        await fs.mkdir(join(reportsDir, key, 'transforms'), { recursive: true })
-
-        return await Promise.all([
-          writeJSON(
-            join(reportsDir, key, 'modules.json'),
-            env.getModulesList(),
-          ),
-          writeJSON(
-            join(reportsDir, key, 'metric-plugins.json'),
-            env.getPluginMetrics(),
-          ),
-          ...Object.entries(env.data.transform)
-            .map(([id, info]) => writeJSON(
-              join(reportsDir, key, 'transforms', `${hash(id)}.json`),
-              <ModuleTransformInfo>{
-                resolvedId: id,
-                transforms: info,
-              },
-            )),
-        ])
-      }),
-  ])
-
-  return targetDir
+    },
+  }
 }
 
 function writeJSON(filename: string, data: any) {

src/node/index.ts

@@ -1,4 +1,4 @@
-import type { Connect, Plugin, ViteDevServer } from 'vite'
+import type { Connect, Plugin, Rollup, ViteDevServer } from 'vite'
 import type { HMRData, RpcFunctions } from '../types'
 import type { InspectContextVite } from './context'
 import type { ViteInspectOptions } from './options'
@@ -8,7 +8,7 @@ import { debounce } from 'perfect-debounce'
 import sirv from 'sirv'
 import { createRPCServer } from 'vite-dev-rpc'
 import { DIR_CLIENT } from '../dirs'
-import { generateBuild } from './build'
+import { createBuildGenerator, createEnvOrderHooks } from './build'
 import { InspectContext } from './context'
 import { hijackPlugin } from './hijack'
 import { createPreviewServer } from './preview'
@@ -39,6 +39,7 @@ export default function PluginInspect(options: ViteInspectOptions = {}): Plugin
   }
 
   const ctx = new InspectContext(options)
+  let onBuildEnd: (envName: string, pluginCtx: Rollup.PluginContext) => Promise<void> | undefined
 
   const timestampRE = /\bt=\d{13}&?\b/
   const trailingSeparatorRE = /[?&]$/
@@ -211,6 +212,24 @@ export default function PluginInspect(options: ViteInspectOptions = {}): Plugin
           return result
         }
       }
+
+      if (build) {
+        const buildGenerator = createBuildGenerator(ctx)
+        onBuildEnd = createEnvOrderHooks<[pluginCtx: Rollup.PluginContext]>(Object.keys(config.environments), {
+          async onFirst() {
+            await buildGenerator.setupOutputDir()
+          },
+          async onEach(pluginCtx) {
+            await buildGenerator.generateForEnv(pluginCtx.environment)
+          },
+          async onLast(pluginCtx) {
+            const dir = buildGenerator.getOutputDir()
+            pluginCtx.environment.logger.info(`${c.green('Inspect report generated at')}  ${c.dim(dir)}`)
+            if (_open && !isCI)
+              createPreviewServer(dir)
+          },
+        })
+      }
     },
     configureServer(server) {
       const rpc = configureServer(server)
@@ -242,14 +261,9 @@ export default function PluginInspect(options: ViteInspectOptions = {}): Plugin
       })
     },
     async buildEnd() {
-      if (!build)
-        return
-      const dir = await generateBuild(ctx)
-
-      this.environment!.logger.info(`${c.green('Inspect report generated at')}  ${c.dim(dir)}`)
-      if (_open && !isCI)
-        createPreviewServer(dir)
+      onBuildEnd?.(this.environment.name, this)
     },
+    sharedDuringBuild: true,
   }
   return plugin
 }