diff --git a/packages/babel/src/traversal/path/introspection.js b/packages/babel/src/traversal/path/introspection.js index 0adafb130b..fbccd768f8 100644 --- a/packages/babel/src/traversal/path/introspection.js +++ b/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 diff --git a/packages/babel/src/traversal/path/replacement.js b/packages/babel/src/traversal/path/replacement.js index abaf3d5778..adee260779 100644 --- a/packages/babel/src/traversal/path/replacement.js +++ b/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; diff --git a/packages/babel/test/path.js b/packages/babel/test/path.js index e203c55c9d..813b2827af 100644 --- a/packages/babel/test/path.js +++ b/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};"); + }); });