feat: add await expression support (#157)

Add `builders.awaitExpression()` and corresponding proxy/type definitions
to allow programmatic creation and parsing of `await` expressions.

Closes #114

Author: Jack-sh1

src/builders.ts

@@ -5,6 +5,7 @@ import { literalToAst } from "./proxy/_utils";
 import type { Proxified } from "./types";
 import { parseExpression } from "./code";
 import { proxifyBinaryExpression } from "./proxy/binary-expression";
+import { proxifyAwaitExpression } from "./proxy/await-expression";
 
 const b = recast.types.builders;
 
@@ -66,6 +67,15 @@ export const builders = {
     );
     return proxifyBinaryExpression(node as any);
   },
+  /**
+  * Create an await expression node.
+  */
+ awaitExpression(argument: any): Proxified {
+  
+   const argAst = argument?.$ast || literalToAst(argument);
+   const node = b.awaitExpression(argAst as any);
+   return proxifyAwaitExpression(node as any);
+ },
   /**
    * Create a proxified version of a literal value.
    */

src/proxy/await-expression.ts

@@ -0,0 +1,22 @@
+import type { AwaitExpression } from "@babel/types";
+import type { ProxifiedAwaitExpression, ProxifiedModule } from "./types";
+import { MagicastError } from "../error";
+import { proxify } from "./proxify";
+import { createProxy } from "./_utils";
+
+export function proxifyAwaitExpression(
+   node: AwaitExpression,
+   mod?: ProxifiedModule,
+ ): ProxifiedAwaitExpression {
+   if (node.type !== "AwaitExpression") {
+     throw new MagicastError("Not an await expression");
+   }
+   return createProxy(
+     node,
+     {
+       $type: "await-expression",
+       $argument: proxify(node.argument, mod),
+     },
+     {},
+   ) as ProxifiedAwaitExpression;
+ }

src/proxy/proxify.ts

@@ -13,6 +13,7 @@ import { proxifyBinaryExpression } from "./binary-expression";
 import { proxifyBlockStatement } from "./block-statement";
 import { proxifyFunctionExpression } from "./function-expression";
 import { LITERALS_AST, LITERALS_TYPEOF } from "./_utils";
+import { proxifyAwaitExpression } from "./await-expression";
 
 const _cache = new WeakMap<ASTNode, any>();
 
@@ -81,6 +82,10 @@ export function proxify<T>(node: ASTNode, mod?: ProxifiedModule): Proxified<T> {
       proxy = proxifyBinaryExpression(node, mod);
       break;
     }
+    case "AwaitExpression":{
+      proxy = proxifyAwaitExpression(node, mod);
+      break;
+    }
     case "BlockStatement": {
       proxy = proxifyBlockStatement(node, mod);
       break;

src/proxy/types.ts

@@ -95,6 +95,11 @@ export type ProxifiedBinaryExpression = ProxyBase & {
   $operator: BinaryOperator;
 };
 
+export type ProxifiedAwaitExpression = ProxyBase & {
+  $type: "await-expression";
+  $argument: Proxified;
+}
+
 export type ProxifiedBlockStatement = ProxyBase & {
   $type: "blockStatement";
   $body: ProxifiedArray;
@@ -171,6 +176,7 @@ export type ProxifiedValue =
   | ProxifiedArrowFunctionExpression
   | ProxifiedFunctionExpression
   | ProxifiedBinaryExpression
+  | ProxifiedAwaitExpression
   | ProxifiedBlockStatement;
 
 export type ProxyType = ProxifiedValue["$type"];

test/builders/await-expression.test.ts

@@ -0,0 +1,41 @@
+ import { describe, expect, it } from "vitest";
+ import { builders, parseModule } from "magicast";
+ import { generate } from "../_utils";
+
+ describe("builders/awaitExpression", () => {
+   it("basic await expression", async () => {
+     const expr = builders.awaitExpression(
+       builders.functionCall("fetch", "https://example.com"),
+     );
+     expect(expr.$type).toBe("await-expression");
+
+     const mod = parseModule("");
+     mod.exports.a = expr;
+
+     expect(await generate(mod)).toMatchInlineSnapshot(`
+       "export const a = await fetch("https://example.com");"
+     `);
+   });
+
+   it("await with raw expression", async () => {
+     const expr = builders.awaitExpression(
+       builders.raw('import("./module")'),
+     );
+     expect(expr.$type).toBe("await-expression");
+
+     const mod = parseModule("");
+     mod.exports.a = expr;
+
+     expect(await generate(mod)).toMatchInlineSnapshot(`
+       "export const a = await import("./module");"
+     `);
+   });
+
+   it("parse existing await expression", async () => {
+     const mod = parseModule('export const a = await fetch("https://example.com");');
+     const a = mod.exports.a;
+     expect(a.$type).toBe("await-expression");
+     expect(a.$argument.$type).toBe("function-call");
+     expect(a.$argument.$callee).toBe("fetch");
+   });
+ });