Merge pull request #2469 from jridgewell/replace-implicit-arrow-with-block

Replace arrow expression body with block statement

packages/babel/src/traversal/path/introspection.js

@@ -106,7 +106,7 @@ export function isNodeType(type: string): boolean {
 }
 
 /**
- * This checks whether or now we're in one of the following positions:
+ * This checks whether or not we're in one of the following positions:
  *
  *   for (KEY in right);
  *   for (KEY;;);
@@ -115,9 +115,31 @@ export function isNodeType(type: string): boolean {
  * to tell the path replacement that it's ok to replace this with an expression.
  */
 
- export function canHaveVariableDeclarationOrExpression() {
-    return (this.key === "init" || this.key === "left") && this.parentPath.isFor();
- }
+export function canHaveVariableDeclarationOrExpression() {
+  return (this.key === "init" || this.key === "left") && this.parentPath.isFor();
+}
+
+/**
+ * This checks whether we are swapping an arrow function's body between an
+ * expression and a block statement (or vice versa).
+ *
+ * This is because arrow functions may implicitly return an expression, which
+ * is the same as containing a block statement.
+ */
+
+export function canSwapBetweenExpressionAndStatement(replacement) {
+  if (this.key !== "body" || !this.parentPath.isArrowFunctionExpression()) {
+    return false;
+  }
+
+  if (this.isExpression()) {
+    return t.isBlockStatement(replacement);
+  } else if (this.isBlockStatement()) {
+    return t.isExpression(replacement);
+  }
+
+  return false;
+}
 
 /**
  * Check whether the current path references a completion record

packages/babel/src/traversal/path/replacement.js

@@ -137,14 +137,18 @@ export function replaceWith(replacement, whateverAllowed) {
     return this.replaceWithSourceString();
   }
 
-  // replacing a statement with an expression so wrap it in an expression statement
-  if (this.isNodeType("Statement") && t.isExpression(replacement) && !this.canHaveVariableDeclarationOrExpression()) {
-    replacement = t.expressionStatement(replacement);
+  if (this.isNodeType("Statement") && t.isExpression(replacement)) {
+    if (!this.canHaveVariableDeclarationOrExpression() && !this.canSwapBetweenExpressionAndStatement(replacement)) {
+      // replacing a statement with an expression so wrap it in an expression statement
+      replacement = t.expressionStatement(replacement);
+    }
   }
 
-  // replacing an expression with a statement so let's explode it
   if (this.isNodeType("Expression") && t.isStatement(replacement)) {
-    return this.replaceExpressionWithStatements([replacement]);
+    if (!this.canSwapBetweenExpressionAndStatement(replacement)) {
+      // replacing an expression with a statement so let's explode it
+      return this.replaceExpressionWithStatements([replacement]);
+    }
   }
 
   var oldNode = this.node;

packages/babel/test/path.js

@@ -20,4 +20,50 @@ suite("traversal path", function () {
 
     chai.expect(actualCode).to.be.equal("console.whatever();");
   });
+
+  test("replaceWith (arrow expression body to block statement body)", function () {
+    var expectCode = "var fn = () => true;";
+
+    var actualCode = transform(expectCode, {
+      blacklist: "strict",
+      plugins: [new Plugin("foobar", {
+        visitor: {
+          ArrowFunctionExpression: function () {
+            this.get("body").replaceWith({
+              type: "BlockStatement",
+              body: [{
+                type: "ReturnStatement",
+                argument: {
+                  type: "Literal",
+                  value: true
+                }
+              }]
+            });
+          }
+        }
+      })]
+    }).code;
+
+    chai.expect(actualCode).to.be.equal("var fn = function fn() {\n  return true;\n};");
+  });
+
+  test("replaceWith (arrow block statement body to expression body)", function () {
+    var expectCode = "var fn = () => { return true; }";
+
+    var actualCode = transform(expectCode, {
+      blacklist: "strict",
+      plugins: [new Plugin("foobar", {
+        visitor: {
+          ArrowFunctionExpression: function () {
+            this.get("body").replaceWith({
+              type: "Literal",
+              value: true
+            });
+          }
+        }
+      })]
+    }).code;
+
+    chai.expect(actualCode).to.be.equal("var fn = function fn() {\n  return true;\n};");
+  });
 });