feat(helpers): Handle Vite config objects declared as variables in addVitePlugin

src/helpers/config.ts

@@ -1,7 +1,69 @@
-import type { ProxifiedModule } from "../types";
+import { Program, VariableDeclarator } from "@babel/types";
+import { generateCode, parseExpression } from "../code";
+import { MagicastError } from "../error";
+import type { Proxified, ProxifiedModule, ProxifiedObject } from "../types";
 
 export function getDefaultExportOptions(magicast: ProxifiedModule<any>) {
-  return magicast.exports.default.$type === "function-call"
-    ? magicast.exports.default.$args[0]
-    : magicast.exports.default;
+  return configFromNode(magicast.exports.default);
+}
+
+/**
+ * Returns the vite config object from a variable declaration thats
+ * exported as the default export.
+ *
+ * Example:
+ *
+ * ```js
+ * const config = {};
+ * export default config;
+ * ```
+ *
+ * @param magicast the module
+ *
+ * @returns an object containing the proxified config object and the
+ *          declaration "parent" to attach the modified config to later.
+ *          If no config declaration is found, undefined is returned.
+ */
+export function getConfigFromVariableDeclaration(
+  magicast: ProxifiedModule<any>
+): {
+  declaration: VariableDeclarator;
+  config: ProxifiedObject<any> | undefined;
+} {
+  if (magicast.exports.default.$type !== "identifier") {
+    throw new MagicastError(
+      `Not supported: Cannot modify this kind of default export (${magicast.exports.default.$type})`
+    );
+  }
+
+  const configDecalarationId = magicast.exports.default.$name;
+
+  for (const node of (magicast.$ast as Program).body) {
+    if (node.type === "VariableDeclaration") {
+      for (const declaration of node.declarations) {
+        if (
+          declaration.id.type === "Identifier" &&
+          declaration.id.name === configDecalarationId &&
+          declaration.init
+        ) {
+          const init = declaration.init;
+
+          const configExpression = parseExpression(generateCode(init).code);
+
+          return {
+            declaration,
+            config: configFromNode(configExpression),
+          };
+        }
+      }
+    }
+  }
+  throw new MagicastError("Couldn't find config declaration");
+}
+
+function configFromNode(node: Proxified<any>): ProxifiedObject<any> {
+  if (node.$type === "function-call") {
+    return node.$args[0];
+  }
+  return node;
 }

src/helpers/vite.ts

@@ -1,6 +1,14 @@
-import type { ProxifiedFunctionCall, ProxifiedModule } from "../proxy/types";
+import type {
+  Proxified,
+  ProxifiedFunctionCall,
+  ProxifiedModule,
+} from "../proxy/types";
 import { builders } from "../builders";
-import { getDefaultExportOptions } from "./config";
+import { generateCode } from "../code";
+import {
+  getConfigFromVariableDeclaration,
+  getDefaultExportOptions,
+} from "./config";
 import { deepMergeObject } from "./deep-merge";
 
 export interface AddVitePluginOptions {
@@ -45,18 +53,13 @@ export function addVitePlugin(
   magicast: ProxifiedModule<any>,
   plugin: AddVitePluginOptions
 ) {
-  const config = getDefaultExportOptions(magicast);
-
-  const insertionIndex = plugin.index ?? config.plugins?.length ?? 0;
+  const config: Proxified | undefined = getDefaultExportOptions(magicast);
 
-  config.plugins ||= [];
-  config.plugins.splice(
-    insertionIndex,
-    0,
-    plugin.options
-      ? builders.functionCall(plugin.constructor, plugin.options)
-      : builders.functionCall(plugin.constructor)
-  );
+  if (config.$type === "identifier") {
+    insertPluginIntoVariableDeclarationConfig(magicast, plugin);
+  } else {
+    insertPluginIntoConfig(plugin, config);
+  }
 
   magicast.imports.$add({
     from: plugin.from,
@@ -106,3 +109,46 @@ export function updateVitePluginConfig(
 
   return true;
 }
+
+/**
+ * Insert @param plugin into a config object that's declared as a variable in
+ * the module (@param magicast).
+ */
+function insertPluginIntoVariableDeclarationConfig(
+  magicast: ProxifiedModule<any>,
+  plugin: AddVitePluginOptions
+) {
+  const { config: configObject, declaration } =
+    getConfigFromVariableDeclaration(magicast);
+
+  insertPluginIntoConfig(plugin, configObject);
+
+  if (declaration.init) {
+    if (declaration.init.type === "ObjectExpression") {
+      // @ts-ignore this works despite the type error because of recast
+      declaration.init = generateCode(configObject).code;
+    } else if (
+      declaration.init.type === "CallExpression" &&
+      declaration.init.callee.type === "Identifier"
+    ) {
+      // @ts-ignore this works despite the type error because of recast
+      declaration.init = generateCode(
+        builders.functionCall(declaration.init.callee.name, configObject)
+      ).code;
+    }
+  }
+}
+
+function insertPluginIntoConfig(plugin: AddVitePluginOptions, config: any) {
+  const insertionIndex = plugin.index ?? config.plugins?.length ?? 0;
+
+  config.plugins ||= [];
+
+  config.plugins.splice(
+    insertionIndex,
+    0,
+    plugin.options
+      ? builders.functionCall(plugin.constructor, plugin.options)
+      : builders.functionCall(plugin.constructor)
+  );
+}

test/helpers/vite.test.ts

@@ -116,4 +116,69 @@ export default defineConfig({
       });"
     `);
   });
+
+  it("handles default export from identifier (fn call)", () => {
+    const code = `
+      import { defineConfig } from 'vite';
+      import { somePlugin1, somePlugin2 } from 'some-module'
+
+      const config = defineConfig({
+        plugins: [somePlugin1(), somePlugin2()]
+      });
+
+      export default config;
+    `;
+
+    const mod = parseModule(code);
+
+    addVitePlugin(mod, {
+      from: "vite-plugin-pwa",
+      imported: "VitePWA",
+      constructor: "VitePWA",
+    });
+
+    expect(generate(mod)).toMatchInlineSnapshot(`
+      "import { VitePWA } from \\"vite-plugin-pwa\\";
+      import { defineConfig } from \\"vite\\";
+      import { somePlugin1, somePlugin2 } from \\"some-module\\";
+
+      const config = defineConfig({
+        plugins: [somePlugin1(), somePlugin2(), VitePWA()],
+      });
+
+      export default config;"
+    `);
+  });
+
+  it("handles default export from identifier (object)", () => {
+    const code = `
+      import { somePlugin1, somePlugin2 } from 'some-module'
+
+      const myConfig = {
+        plugins: [somePlugin1(), somePlugin2()]
+      };
+
+      export default myConfig;
+    `;
+
+    const mod = parseModule(code);
+
+    addVitePlugin(mod, {
+      index: 1,
+      from: "vite-plugin-pwa",
+      imported: "VitePWA",
+      constructor: "VitePWA",
+    });
+
+    expect(generate(mod)).toMatchInlineSnapshot(`
+      "import { VitePWA } from \\"vite-plugin-pwa\\";
+      import { somePlugin1, somePlugin2 } from \\"some-module\\";
+
+      const myConfig = {
+        plugins: [somePlugin1(), VitePWA(), somePlugin2()],
+      };
+
+      export default myConfig;"
+    `);
+  });
 });