feat: imports support (#11)

src/proxy/_utils.ts

@@ -5,6 +5,8 @@ import { proxifyFunctionCall } from "./function-call";
 import { proxifyObject } from "./object";
 import { ProxyUtils, Proxified } from "./types";
 
+const b = recast.types.builders;
+
 const literalTypes = new Set([
   "Literal",
   "StringLiteral",
@@ -64,46 +66,42 @@ const PROXY_KEY = "__magicast_proxy";
 
 export function literalToAst(value: any): ESNode {
   if (value === undefined) {
-    return recast.types.builders.identifier("undefined") as any;
+    return b.identifier("undefined") as any;
   }
   if (value === null) {
     // eslint-disable-next-line unicorn/no-null
-    return recast.types.builders.literal(null) as any;
+    return b.literal(null) as any;
   }
   if (Array.isArray(value)) {
-    return recast.types.builders.arrayExpression(
-      value.map((n) => literalToAst(n)) as any
-    ) as any;
+    return b.arrayExpression(value.map((n) => literalToAst(n)) as any) as any;
   }
   if (typeof value === "object") {
     if (PROXY_KEY in value) {
       return value.$ast;
     }
-    return recast.types.builders.objectExpression(
+    return b.objectExpression(
       Object.entries(value).map(([key, value]) => {
-        return recast.types.builders.property(
+        return b.property(
           "init",
-          recast.types.builders.identifier(key),
+          b.identifier(key),
           literalToAst(value) as any
         ) as any;
       })
     ) as any;
   }
-  return recast.types.builders.literal(value) as any;
+  return b.literal(value) as any;
 }
 
 export function makeProxyUtils<T extends object>(
   node: ESNode,
   extend: T = {} as T
 ): ProxyUtils & T {
-  return {
-    [PROXY_KEY]: true,
-    get $ast() {
-      return node;
-    },
-    $type: "object",
-    ...extend,
-  } as ProxyUtils & T;
+  const obj = extend as ProxyUtils & T;
+  // @ts-expect-error internal property
+  obj[PROXY_KEY] = true;
+  obj.$ast = node;
+  obj.$type ||= "object";
+  return obj;
 }
 
 export function createProxy<T extends object>(
@@ -124,6 +122,16 @@ export function createProxy<T extends object>(
           return handler.get(target, key, receiver);
         }
       },
+      set(target: T, key: string | symbol, value: any, receiver: any) {
+        if (key in utils) {
+          (utils as any)[key] = value;
+          return true;
+        }
+        if (handler.set) {
+          return handler.set(target, key, value, receiver);
+        }
+        return false;
+      },
     }
   ) as Proxified<T>;
 }

src/proxy/exports.ts

@@ -0,0 +1,101 @@
+import * as recast from "recast";
+import { Program } from "@babel/types";
+import { createProxy, literalToAst, proxify } from "./_utils";
+
+const b = recast.types.builders;
+
+export function createExportsProxy(root: Program) {
+  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"
+        ? b.exportDefaultDeclaration(node)
+        : (b.exportNamedDeclaration(
+            b.variableDeclaration("const", [
+              b.variableDeclarator(b.identifier(key), node),
+            ])
+          ) as any)
+    );
+  };
+
+  return createProxy(
+    root,
+    {
+      $type: "exports",
+    },
+    {
+      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;
+      },
+    }
+  );
+}

src/proxy/imports.ts

@@ -0,0 +1,206 @@
+/* eslint-disable unicorn/no-nested-ternary */
+import * as recast from "recast";
+import {
+  ImportDeclaration,
+  ImportDefaultSpecifier,
+  ImportNamespaceSpecifier,
+  ImportSpecifier,
+  Program,
+} from "@babel/types";
+import { createProxy } from "./_utils";
+import {
+  ImportItemInput,
+  ProxifiedImportItem,
+  ProxifiedImportsMap,
+} from "./types";
+
+const b = recast.types.builders;
+const _importProxyCache = new WeakMap<any, ProxifiedImportItem>();
+
+export function creatImportProxy(
+  node: ImportDeclaration,
+  specifier:
+    | ImportSpecifier
+    | ImportNamespaceSpecifier
+    | ImportDefaultSpecifier,
+  root: Program
+): ProxifiedImportItem {
+  if (_importProxyCache.has(specifier)) {
+    return _importProxyCache.get(specifier)!;
+  }
+  const proxy = createProxy(
+    specifier,
+    {
+      get $declaration() {
+        return node;
+      },
+      get imported() {
+        if (specifier.type === "ImportDefaultSpecifier") {
+          return "default";
+        }
+        if (specifier.type === "ImportNamespaceSpecifier") {
+          return "*";
+        }
+        if (specifier.imported.type === "Identifier") {
+          return specifier.imported.name;
+        }
+        return specifier.imported.value;
+      },
+      set imported(value) {
+        if (specifier.type !== "ImportSpecifier") {
+          throw new Error("Changing import name is not yet implemented");
+        }
+        if (specifier.imported.type === "Identifier") {
+          specifier.imported.name = value;
+        } else {
+          specifier.imported.value = value;
+        }
+      },
+      get local() {
+        return specifier.local.name;
+      },
+      set local(value) {
+        specifier.local.name = value;
+      },
+      get from() {
+        return node.source.value;
+      },
+      set from(value) {
+        if (value === node.source.value) {
+          return;
+        }
+
+        node.specifiers = node.specifiers.filter((s) => s !== specifier);
+        if (node.specifiers.length === 0) {
+          root.body = root.body.filter((s) => s !== node);
+        }
+
+        const declaration = root.body.find(
+          (i) => i.type === "ImportDeclaration" && i.source.value === value
+        ) as ImportDeclaration | undefined;
+        if (!declaration) {
+          root.body.unshift(
+            b.importDeclaration(
+              [specifier as any],
+              b.stringLiteral(value)
+            ) as any
+          );
+        } else {
+          // TODO: insert after the last import maybe?
+          declaration.specifiers.push(specifier as any);
+        }
+      },
+      toJSON() {
+        return {
+          imported: this.imported,
+          local: this.local,
+          from: this.from,
+        };
+      },
+    },
+    {
+      ownKeys() {
+        return ["imported", "local", "from"];
+      },
+    }
+  ) as ProxifiedImportItem;
+  _importProxyCache.set(specifier, proxy);
+  return proxy;
+}
+
+export function createImportsProxy(root: Program) {
+  // TODO: cache
+  const getAllImports = () => {
+    const imports: ReturnType<typeof creatImportProxy>[] = [];
+    for (const n of root.body) {
+      if (n.type === "ImportDeclaration") {
+        for (const specifier of n.specifiers) {
+          imports.push(creatImportProxy(n, specifier, root));
+        }
+      }
+    }
+    return imports;
+  };
+
+  const updateImport = (key: string, value: ImportItemInput) => {
+    const imports = getAllImports();
+    const item = imports.find((i) => i.local === key);
+    const local = value.local || key;
+    if (item) {
+      item.imported = value.imported;
+      item.local = local;
+      item.from = value.from;
+      return true;
+    }
+
+    const specifier =
+      value.imported === "default"
+        ? b.importDefaultSpecifier(b.identifier(local))
+        : value.imported === "*"
+        ? b.importNamespaceSpecifier(b.identifier(local))
+        : b.importSpecifier(b.identifier(value.imported), b.identifier(local));
+
+    const declaration = imports.find(
+      (i) => i.from === value.from
+    )?.$declaration;
+    if (!declaration) {
+      root.body.unshift(
+        b.importDeclaration([specifier], b.stringLiteral(value.from)) as any
+      );
+    } else {
+      // TODO: insert after the last import maybe?
+      declaration.specifiers.push(specifier as any);
+    }
+    return true;
+  };
+
+  const removeImport = (key: string) => {
+    const item = getAllImports().find((i) => i.local === key);
+    if (!item) {
+      return false;
+    }
+    const node = item.$declaration;
+    const specifier = item.$ast;
+    node.specifiers = node.specifiers.filter((s) => s !== specifier);
+    if (node.specifiers.length === 0) {
+      root.body = root.body.filter((n) => n !== node);
+    }
+    return true;
+  };
+
+  const proxy = createProxy(
+    root,
+    {
+      $type: "imports",
+      $add(item: ImportItemInput) {
+        proxy[item.local || item.imported] = item as any;
+      },
+      toJSON() {
+        // eslint-disable-next-line unicorn/no-array-reduce
+        return getAllImports().reduce((acc, i) => {
+          acc[i.local] = i;
+          return acc;
+        }, {} as any);
+      },
+    },
+    {
+      get(_, prop) {
+        return getAllImports().find((i) => i.local === prop);
+      },
+      set(_, prop, value) {
+        return updateImport(prop as string, value);
+      },
+      deleteProperty(_, prop) {
+        return removeImport(prop as string);
+      },
+      ownKeys() {
+        return getAllImports().map((i) => i.local);
+      },
+      has(_, prop) {
+        return getAllImports().some((i) => i.local === prop);
+      },
+    }
+  ) as any as ProxifiedImportsMap;
+
+  return proxy;
+}