feat: use proxy for top level module (#8)

* feat: able to manipulate exports

* feat: add delete operation

* chore: lint

Author: antfu

src/code.ts

@@ -1,31 +1,35 @@
 import { promises as fsp } from "node:fs";
 import { print, parse, Options as ParseOptions } from "recast";
-import { ModuleNode } from "./module";
 import { getBabelParser } from "./babel";
-import { ESNode, ParsedFileNode } from "./types";
+import { ESNode, ParsedFileNode, ProxifiedModule } from "./types";
+import { proxifyModule } from "./proxy/module";
 
-export function parseCode(code: string, options?: ParseOptions): ModuleNode {
+export function parseCode<T = any>(
+  code: string,
+  options?: ParseOptions
+): ProxifiedModule<T> {
   const node: ParsedFileNode = parse(code, {
     parser: options?.parser || getBabelParser(),
     ...options,
   });
-  return new ModuleNode(node);
+  return proxifyModule(node);
 }
 
 export function generateCode(
-  node: { ast: ESNode } | ESNode,
+  node: { $ast: ESNode } | ESNode,
   options?: ParseOptions
 ): { code: string; map?: any } {
-  const { code, map } = print("ast" in node ? node.ast : node, {
+  const ast = "$ast" in node ? node.$ast : node;
+  const { code, map } = print(ast, {
     ...options,
   });
   return { code, map };
 }
 
-export async function loadFile(
+export async function loadFile<T = any>(
   filename: string,
   options: ParseOptions = {}
-): Promise<ModuleNode> {
+): Promise<ProxifiedModule<T>> {
   const contents = await fsp.readFile(filename, "utf8");
   options.sourceFileName = options.sourceFileName ?? filename;
   return parseCode(contents, options);

src/index.ts

@@ -1,3 +1,2 @@
 export * from "./code";
-export * from "./module";
 export * from "./types";

src/module.ts

@@ -0,0 +1,115 @@
+import * as recast from "recast";
+import { ESNode, ParsedFileNode } from "../types";
+import { ProxifiedModule } from "./types";
+import { createProxy, literalToAst, proxify } from "./_utils";
+
+export function proxifyModule<T>(ast: ParsedFileNode): ProxifiedModule<T> {
+  const root = ast.program as ESNode;
+  if (root.type !== "Program") {
+    throw new Error(`Cannot proxify ${ast.type} as module`);
+  }
+
+  const findExport = (key: string) => {
+    const type =
+      key === "default" ? "ExportDefaultDeclaration" : "ExportNamedDeclaration";
+
+    for (const n of root.body) {
+      if (n.type === type) {
+        if (key === "default") {
+          return n.declaration;
+        }
+        if (n.declaration && "declarations" in n.declaration) {
+          const dec = n.declaration.declarations[0];
+          if ("name" in dec.id && dec.id.name === key) {
+            return dec.init as any;
+          }
+        }
+      }
+    }
+  };
+
+  const updateOrAddExport = (key: string, value: any) => {
+    const type =
+      key === "default" ? "ExportDefaultDeclaration" : "ExportNamedDeclaration";
+
+    const node = literalToAst(value) as any;
+    for (const n of root.body) {
+      if (n.type === type) {
+        if (key === "default") {
+          n.declaration = node;
+          return;
+        }
+        if (n.declaration && "declarations" in n.declaration) {
+          const dec = n.declaration.declarations[0];
+          if ("name" in dec.id && dec.id.name === key) {
+            dec.init = node;
+            return;
+          }
+        }
+      }
+    }
+
+    root.body.push(
+      key === "default"
+        ? recast.types.builders.exportDefaultDeclaration(node)
+        : (recast.types.builders.exportNamedDeclaration(
+            recast.types.builders.variableDeclaration("const", [
+              recast.types.builders.variableDeclarator(
+                recast.types.builders.identifier(key),
+                node
+              ),
+            ])
+          ) as any)
+    );
+  };
+
+  const exportsProxy = createProxy(
+    root,
+    {},
+    {
+      get(_, prop) {
+        const node = findExport(prop as string);
+        if (node) {
+          return proxify(node);
+        }
+      },
+      set(_, prop, value) {
+        updateOrAddExport(prop as string, value);
+        return true;
+      },
+      deleteProperty(_, prop) {
+        const type =
+          prop === "default"
+            ? "ExportDefaultDeclaration"
+            : "ExportNamedDeclaration";
+
+        for (let i = 0; i < root.body.length; i++) {
+          const n = root.body[i];
+          if (n.type === type) {
+            if (prop === "default") {
+              root.body.splice(i, 1);
+              return true;
+            }
+            if (n.declaration && "declarations" in n.declaration) {
+              const dec = n.declaration.declarations[0];
+              if ("name" in dec.id && dec.id.name === prop) {
+                root.body.splice(i, 1);
+                return true;
+              }
+            }
+          }
+        }
+        return false;
+      },
+    }
+  );
+
+  return {
+    $ast: root,
+    $type: "module",
+    exports: exportsProxy,
+    get imports() {
+      throw new Error("Not implemented");
+    },
+  } as any;
+}

src/proxy/module.ts

@@ -35,3 +35,8 @@ export type Proxified<T = any> = T extends
       [K in keyof T]: Proxified<T[K]>;
     } & ProxyUtils
   : T;
+
+export interface ProxifiedModule<T = Record<string, unknown>> {
+  exports: Proxified<T>;
+  imports: unknown[];
+}

src/proxy/types.ts

@@ -205,4 +205,49 @@ describe("magicast", () => {
       }
     `);
   });
+
+  it("manipulate exports", () => {
+    const mod = parseCode("");
+
+    expect(mod.exports).toMatchInlineSnapshot(`{}`);
+    expect(generate(mod)).toMatchInlineSnapshot('""');
+
+    mod.exports.default = { foo: "1" };
+
+    expect(generate(mod)).toMatchInlineSnapshot(`
+      "export default {
+        foo: \\"1\\",
+      };"
+    `);
+
+    mod.exports.default.foo = 2;
+
+    expect(generate(mod)).toMatchInlineSnapshot(`
+      "export default {
+        foo: 2,
+      };"
+    `);
+
+    mod.exports.named ||= [];
+    mod.exports.named.push("a");
+
+    expect(generate(mod)).toMatchInlineSnapshot(`
+      "export default {
+        foo: 2,
+      };
+
+      export const named = [\\"a\\"];"
+    `);
+
+    // delete
+    delete mod.exports.default;
+
+    expect(generate(mod)).toMatchInlineSnapshot(
+      '"export const named = [\\"a\\"];"'
+    );
+
+    delete mod.exports.named;
+
+    expect(generate(mod)).toMatchInlineSnapshot('""');
+  });
 });