Transform await in computed class keys

packages/babel-helper-remap-async-to-generator/package.json

@@ -15,12 +15,16 @@
   "main": "./lib/index.js",
   "dependencies": {
     "@babel/helper-annotate-as-pure": "workspace:^",
+    "@babel/helper-environment-visitor": "workspace:^",
     "@babel/helper-wrap-function": "workspace:^",
     "@babel/types": "workspace:^"
   },
   "devDependencies": {
     "@babel/traverse": "workspace:^"
   },
+  "peerDependencies": {
+    "@babel/core": "workspace:^"
+  },
   "engines": {
     "node": ">=6.9.0"
   },

packages/babel-helper-remap-async-to-generator/src/index.ts

@@ -1,8 +1,10 @@
 /* @noflow */
 
-import type { NodePath, Visitor } from "@babel/traverse";
+import type { NodePath } from "@babel/traverse";
 import wrapFunction from "@babel/helper-wrap-function";
 import annotateAsPure from "@babel/helper-annotate-as-pure";
+import environmentVisitor from "@babel/helper-environment-visitor";
+import { traverse } from "@babel/core";
 import {
   callExpression,
   cloneNode,
@@ -12,23 +14,26 @@ import {
 } from "@babel/types";
 import type * as t from "@babel/types";
 
-const awaitVisitor: Visitor<{ wrapAwait: t.Expression }> = {
-  Function(path) {
-    path.skip();
-  },
+const awaitVisitor = traverse.visitors.merge<{ wrapAwait: t.Expression }>([
+  {
+    ArrowFunctionExpression(path) {
+      path.skip();
+    },
 
-  AwaitExpression(path, { wrapAwait }) {
-    const argument = path.get("argument");
+    AwaitExpression(path, { wrapAwait }) {
+      const argument = path.get("argument");
 
-    path.replaceWith(
-      yieldExpression(
-        wrapAwait
-          ? callExpression(cloneNode(wrapAwait), [argument.node])
-          : argument.node,
-      ),
-    );
+      path.replaceWith(
+        yieldExpression(
+          wrapAwait
+            ? callExpression(cloneNode(wrapAwait), [argument.node])
+            : argument.node,
+        ),
+      );
+    },
   },
-};
+  environmentVisitor,
+]);
 
 export default function (
   path: NodePath<any>,

packages/babel-plugin-proposal-async-generator-functions/package.json

@@ -17,6 +17,7 @@
     "babel-plugin"
   ],
   "dependencies": {
+    "@babel/helper-environment-visitor": "workspace:^",
     "@babel/helper-plugin-utils": "workspace:^",
     "@babel/helper-remap-async-to-generator": "workspace:^",
     "@babel/plugin-syntax-async-generators": "^7.8.4"

packages/babel-plugin-proposal-async-generator-functions/src/index.ts

@@ -1,66 +1,72 @@
 import { declare } from "@babel/helper-plugin-utils";
 import remapAsyncToGenerator from "@babel/helper-remap-async-to-generator";
 import syntaxAsyncGenerators from "@babel/plugin-syntax-async-generators";
-import { types as t } from "@babel/core";
-import type { PluginPass } from "@babel/core";
 import type { NodePath, Visitor } from "@babel/traverse";
+import { traverse, types as t, type PluginPass } from "@babel/core";
 import rewriteForAwait from "./for-await";
+import environmentVisitor from "@babel/helper-environment-visitor";
 
 export default declare(api => {
   api.assertVersion(7);
 
-  const yieldStarVisitor: Visitor<PluginPass> = {
-    Function(path) {
-      path.skip();
-    },
+  const yieldStarVisitor = traverse.visitors.merge<PluginPass>([
+    {
+      ArrowFunctionExpression(path) {
+        path.skip();
+      },
 
-    YieldExpression({ node }, state) {
-      if (!node.delegate) return;
-      const callee = state.addHelper("asyncGeneratorDelegate");
-      node.argument = t.callExpression(callee, [
-        t.callExpression(state.addHelper("asyncIterator"), [node.argument]),
-        state.addHelper("awaitAsyncGenerator"),
-      ]);
+      YieldExpression({ node }, state) {
+        if (!node.delegate) return;
+        const callee = state.addHelper("asyncGeneratorDelegate");
+        node.argument = t.callExpression(callee, [
+          t.callExpression(state.addHelper("asyncIterator"), [node.argument]),
+          state.addHelper("awaitAsyncGenerator"),
+        ]);
+      },
     },
-  };
+    environmentVisitor,
+  ]);
 
-  const forAwaitVisitor: Visitor<PluginPass> = {
-    Function(path) {
-      path.skip();
-    },
+  const forAwaitVisitor = traverse.visitors.merge<PluginPass>([
+    {
+      ArrowFunctionExpression(path) {
+        path.skip();
+      },
 
-    ForOfStatement(path: NodePath<t.ForOfStatement>, { file }) {
-      const { node } = path;
-      if (!node.await) return;
+      ForOfStatement(path: NodePath<t.ForOfStatement>, { file }) {
+        const { node } = path;
+        if (!node.await) return;
 
-      const build = rewriteForAwait(path, {
-        getAsyncIterator: file.addHelper("asyncIterator"),
-      });
+        const build = rewriteForAwait(path, {
+          getAsyncIterator: file.addHelper("asyncIterator"),
+        });
 
-      const { declar, loop } = build;
-      const block = loop.body as t.BlockStatement;
+        const { declar, loop } = build;
+        const block = loop.body as t.BlockStatement;
 
-      // ensure that it's a block so we can take all its statements
-      path.ensureBlock();
+        // ensure that it's a block so we can take all its statements
+        path.ensureBlock();
 
-      // add the value declaration to the new loop body
-      if (declar) {
-        block.body.push(declar);
-      }
+        // add the value declaration to the new loop body
+        if (declar) {
+          block.body.push(declar);
+        }
 
-      // push the rest of the original loop body onto our new body
-      block.body.push(...path.node.body.body);
+        // push the rest of the original loop body onto our new body
+        block.body.push(...path.node.body.body);
 
-      t.inherits(loop, node);
-      t.inherits(loop.body, node.body);
+        t.inherits(loop, node);
+        t.inherits(loop.body, node.body);
 
-      if (build.replaceParent) {
-        path.parentPath.replaceWithMultiple(build.node);
-      } else {
-        path.replaceWithMultiple(build.node);
-      }
+        if (build.replaceParent) {
+          path.parentPath.replaceWithMultiple(build.node);
+        } else {
+          path.replaceWithMultiple(build.node);
+        }
+      },
     },
-  };
+    environmentVisitor,
+  ]);
 
   const visitor: Visitor<PluginPass> = {
     Function(path, state) {

packages/babel-plugin-proposal-async-generator-functions/test/fixtures/async-generators/transform-class-keys/input.js

@@ -0,0 +1,9 @@
+async function* fn() {
+    class A {
+      [yield 1]() {}
+    }
+
+    class B extends A {
+      [await 1]() {}
+    }
+}

packages/babel-plugin-proposal-async-generator-functions/test/fixtures/async-generators/transform-class-keys/options.json

@@ -0,0 +1,3 @@
+{
+    "plugins": ["proposal-async-generator-functions"]
+}

packages/babel-plugin-proposal-async-generator-functions/test/fixtures/async-generators/transform-class-keys/output.js

@@ -0,0 +1,18 @@
+function fn() {
+  return _fn.apply(this, arguments);
+}
+
+function _fn() {
+  _fn = babelHelpers.wrapAsyncGenerator(function* () {
+    class A {
+      [yield 1]() {}
+
+    }
+
+    class B extends A {
+      [yield babelHelpers.awaitAsyncGenerator(1)]() {}
+
+    }
+  });
+  return _fn.apply(this, arguments);
+}

yarn.lock

@@ -845,9 +845,12 @@ __metadata:
   resolution: "@babel/helper-remap-async-to-generator@workspace:packages/babel-helper-remap-async-to-generator"
   dependencies:
     "@babel/helper-annotate-as-pure": "workspace:^"
+    "@babel/helper-environment-visitor": "workspace:^"
     "@babel/helper-wrap-function": "workspace:^"
     "@babel/traverse": "workspace:^"
     "@babel/types": "workspace:^"
+  peerDependencies:
+    "@babel/core": "workspace:^"
   languageName: unknown
   linkType: soft
 
@@ -1197,6 +1200,7 @@ __metadata:
   resolution: "@babel/plugin-proposal-async-generator-functions@workspace:packages/babel-plugin-proposal-async-generator-functions"
   dependencies:
     "@babel/core": "workspace:^"
+    "@babel/helper-environment-visitor": "workspace:^"
     "@babel/helper-plugin-test-runner": "workspace:^"
     "@babel/helper-plugin-utils": "workspace:^"
     "@babel/helper-remap-async-to-generator": "workspace:^"