feat: environment in hooks, context vs param (#16261)

Author: patak-dev

packages/vite/rollup.dts.config.ts

@@ -48,6 +48,8 @@ const identifierWithTrailingDollarRE = /\b(\w+)\$\d+\b/g
 const identifierReplacements: Record<string, Record<string, string>> = {
   rollup: {
     Plugin$1: 'rollup.Plugin',
+    PluginContext$1: 'rollup.PluginContext',
+    TransformPluginContext$1: 'rollup.TransformPluginContext',
     TransformResult$2: 'rollup.TransformResult',
   },
   esbuild: {

packages/vite/src/node/build.ts

@@ -14,6 +14,7 @@ import type {
   RollupLog,
   RollupOptions,
   RollupOutput,
+  PluginContext as RollupPluginContext,
   RollupWatcher,
   WatcherOptions,
 } from 'rollup'
@@ -1040,10 +1041,10 @@ function wrapEnvironmentResolveId(
   const fn = getHookHandler(hook)
   const handler: Plugin['resolveId'] = function (id, importer, options) {
     return fn.call(
-      this,
+      injectEnvironmentInContext(this),
       id,
       importer,
-      injectEnvironmentFlag(options, environment),
+      injectSsrFlag(options, environment),
     )
   }
 
@@ -1065,8 +1066,12 @@ function wrapEnvironmentLoad(
 
   const fn = getHookHandler(hook)
   const handler: Plugin['load'] = function (id, ...args) {
-    // @ts-expect-error: Receiving options param to be future-proof if Rollup adds it
-    return fn.call(this, id, injectEnvironmentFlag(args[0], environment))
+    return fn.call(
+      injectEnvironmentInContext(this, environment),
+      id,
+      // @ts-expect-error: Receiving options param to be future-proof if Rollup adds it
+      injectSsrFlag(args[0], environment),
+    )
   }
 
   if ('handler' in hook) {
@@ -1088,11 +1093,11 @@ function wrapEnvironmentTransform(
   const fn = getHookHandler(hook)
   const handler: Plugin['transform'] = function (code, importer, ...args) {
     return fn.call(
-      this,
+      injectEnvironmentInContext(this, environment),
       code,
       importer,
       // @ts-expect-error: Receiving options param to be future-proof if Rollup adds it
-      injectEnvironmentFlag(args[0], environment),
+      injectSsrFlag(args[0], environment),
     )
   }
 
@@ -1106,14 +1111,27 @@ function wrapEnvironmentTransform(
   }
 }
 
-function injectEnvironmentFlag<T extends Record<string, any>>(
+function injectEnvironmentInContext(
+  context: RollupPluginContext,
+  environment?: BuildEnvironment,
+) {
+  return new Proxy(context, {
+    get(target, prop, receiver) {
+      if (prop === 'environment') {
+        return environment
+      }
+      return Reflect.get(target, prop, receiver)
+    },
+  })
+}
+
+function injectSsrFlag<T extends Record<string, any>>(
   options?: T,
   environment?: BuildEnvironment,
-): T & { ssr?: boolean; environment?: BuildEnvironment } {
+): T & { ssr?: boolean } {
   const ssr = environment ? environment.name !== 'client' : true
-  return { ...(options ?? {}), ssr, environment } as T & {
+  return { ...(options ?? {}), ssr } as T & {
     ssr?: boolean
-    environment?: BuildEnvironment
   }
 }
 

packages/vite/src/node/plugin.ts

@@ -2,13 +2,12 @@ import type {
   CustomPluginOptions,
   LoadResult,
   ObjectHook,
-  PluginContext,
   ResolveIdResult,
   Plugin as RollupPlugin,
-  TransformPluginContext,
+  PluginContext as RollupPluginContext,
+  TransformPluginContext as RollupTransformPluginContext,
   TransformResult,
 } from 'rollup'
-export type { PluginContext } from 'rollup'
 import type {
   ConfigEnv,
   EnvironmentConfig,
@@ -44,7 +43,19 @@ import type { BuildEnvironment } from './build'
  *
  * If a plugin should be applied only for server or build, a function format
  * config file can be used to conditional determine the plugins to use.
+ *
+ * The current module environment can be accessed from the context for the
+ * buildStart, resolveId, transform, load, and buildEnd, hooks
  */
+
+export interface PluginContext extends RollupPluginContext {
+  environment?: DevEnvironment | BuildEnvironment
+}
+
+export interface TransformPluginContext extends RollupTransformPluginContext {
+  environment?: DevEnvironment | BuildEnvironment
+}
+
 export interface Plugin<A = any> extends RollupPlugin<A> {
   /**
    * Enforce plugin invocation tier similar to webpack loaders.
@@ -198,8 +209,10 @@ export interface Plugin<A = any> extends RollupPlugin<A> {
       options: {
         attributes: Record<string, string>
         custom?: CustomPluginOptions
+        /**
+         * @deprecated use this.environment
+         */
         ssr?: boolean
-        environment?: DevEnvironment | BuildEnvironment
         /**
          * @internal
          */
@@ -213,8 +226,10 @@ export interface Plugin<A = any> extends RollupPlugin<A> {
       this: PluginContext,
       id: string,
       options?: {
+        /**
+         * @deprecated use this.environment
+         */
         ssr?: boolean
-        environment?: DevEnvironment | BuildEnvironment
       },
     ) => Promise<LoadResult> | LoadResult
   >
@@ -224,8 +239,10 @@ export interface Plugin<A = any> extends RollupPlugin<A> {
       code: string,
       id: string,
       options?: {
+        /**
+         * @deprecated use this.environment
+         */
         ssr?: boolean
-        environment?: DevEnvironment | BuildEnvironment
       },
     ) => Promise<TransformResult> | TransformResult
   >

packages/vite/src/node/plugins/asset.ts

@@ -200,7 +200,7 @@ export function assetPlugin(config: ResolvedConfig): Plugin {
 
       // Inherit HMR timestamp if this asset was invalidated
       if (server) {
-        const environment = options?.environment as DevEnvironment | undefined
+        const environment = this.environment as DevEnvironment | undefined
         const mod = environment?.moduleGraph.getModuleById(id)
         if (mod && mod.lastHMRTimestamp > 0) {
           url = injectQuery(url, `t=${mod.lastHMRTimestamp}`)

packages/vite/src/node/plugins/css.ts

@@ -941,7 +941,7 @@ export function cssAnalysisPlugin(config: ResolvedConfig): Plugin {
         return
       }
 
-      const environment = options?.environment
+      const environment = this.environment
       const moduleGraph =
         environment?.mode === 'dev' ? environment.moduleGraph : undefined
       const thisModule = moduleGraph?.getModuleById(id)

packages/vite/src/node/plugins/importAnalysis.ts

@@ -215,7 +215,7 @@ export function importAnalysisPlugin(config: ResolvedConfig): Plugin {
       // running src/node/server/__tests__/pluginContainer.spec.ts
 
       const ssr = options?.ssr === true
-      const environment = (options?.environment as DevEnvironment) || undefined
+      const environment = this.environment as DevEnvironment | undefined
 
       if (!server || !environment) {
         return null

packages/vite/src/node/plugins/worker.ts

@@ -258,7 +258,7 @@ export function webWorkerPlugin(config: ResolvedConfig): Plugin {
           } else if (server) {
             // dynamic worker type we can't know how import the env
             // so we copy /@vite/env code of server transform result into file header
-            const environment = options?.environment
+            const environment = this.environment
             const moduleGraph =
               environment?.mode === 'dev' ? environment.moduleGraph : undefined
             const module = moduleGraph?.getModuleById(ENV_ENTRY)

packages/vite/src/node/server/pluginContainer.ts

@@ -731,7 +731,6 @@ export async function createPluginContainer(
             custom: options?.custom,
             isEntry: !!options?.isEntry,
             ssr,
-            environment,
             scan,
           }),
         )
@@ -786,7 +785,7 @@ export async function createPluginContainer(
         ctx._activePlugin = plugin
         const handler = getHookHandler(plugin.load)
         const result = await handleHookPromise(
-          handler.call(ctx as any, id, { ssr, environment }),
+          handler.call(ctx as any, id, { ssr }),
         )
         if (result != null) {
           if (isObject(result)) {
@@ -822,7 +821,7 @@ export async function createPluginContainer(
         const handler = getHookHandler(plugin.transform)
         try {
           result = await handleHookPromise(
-            handler.call(ctx as any, code, id, { ssr, environment }),
+            handler.call(ctx as any, code, id, { ssr }),
           )
         } catch (e) {
           ctx.error(e)