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

Replace arrow expression body with block statement

packages/babel-traverse/src/path/introspection.js

@@ -116,7 +116,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;;);
@@ -129,6 +129,28 @@ 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-traverse/src/path/replacement.js

@@ -121,14 +121,18 @@ export function replaceWith(replacement) {
     throw new Error("Don't use `path.replaceWith()` with a source string, use `path.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.canHaveVariableDeclarationOrExpression() && !this.canSwapBetweenExpressionAndStatement(replacement)) {
+      // replacing an expression with a statement so let's explode it
+      return this.replaceExpressionWithStatements([replacement]);
+    }
   }
 
   let oldNode = this.node;