From de1431e8c6bb59910c7a65fe6f58b1cfbb5c0ebc Mon Sep 17 00:00:00 2001 From: Justin Ridgewell Date: Sat, 20 Feb 2016 04:36:40 -0500 Subject: [PATCH] Replace arrow expression body with block statement Original PR: https://github.com/babel/babel/pull/2469. Seems this got lost in the v6 changes. - - - Without this, the only way to replace the arrow function is to either manually override its `node.body`, or duplicate the arrow: ```js // Old ArrowFunctionExpression: function (node) { node.body = t.blockStatement(...); // Or return t.ArrowFunctionExpression( node.params, t.blockStatement(...), node.async ); } // New ArrowFunctionExpression: function() { this.get("body").replaceWith(t.blockStatement(...)); } ``` --- packages/babel-core/test/path.js | 134 ++++++++++++++++++ .../babel-traverse/src/path/introspection.js | 24 +++- .../babel-traverse/src/path/replacement.js | 14 +- 3 files changed, 166 insertions(+), 6 deletions(-) diff --git a/packages/babel-core/test/path.js b/packages/babel-core/test/path.js index f56c88049b..1b8ae9044a 100644 --- a/packages/babel-core/test/path.js +++ b/packages/babel-core/test/path.js @@ -19,4 +19,138 @@ 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, { + plugins: [new Plugin({ + visitor: { + ArrowFunctionExpression: function (path) { + path.get("body").replaceWith({ + type: "BlockStatement", + body: [{ + type: "ReturnStatement", + argument: { + type: "BooleanLiteral", + value: true + } + }] + }); + } + } + })] + }).code; + + chai.expect(actualCode).to.be.equal("var 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, { + plugins: [new Plugin({ + visitor: { + ArrowFunctionExpression: function (path) { + path.get("body").replaceWith({ + type: "BooleanLiteral", + value: true + }); + } + } + })] + }).code; + + chai.expect(actualCode).to.be.equal("var fn = () => true;"); + }); + + test("replaceWith (for-in left expression to variable declaration)", function () { + var expectCode = "for (KEY in right);"; + + var actualCode = transform(expectCode, { + plugins: [new Plugin({ + visitor: { + ForInStatement: function (path) { + path.get("left").replaceWith({ + type: "VariableDeclaration", + kind: "var", + declarations: [{ + type: "VariableDeclarator", + id: { + type: "Identifier", + name: "KEY" + } + }] + }); + } + } + })] + }).code; + + chai.expect(actualCode).to.be.equal("for (var KEY in right);"); + }); + + test("replaceWith (for-in left variable declaration to expression)", function () { + var expectCode = "for (var KEY in right);"; + + var actualCode = transform(expectCode, { + plugins: [new Plugin({ + visitor: { + ForInStatement: function (path) { + path.get("left").replaceWith({ + type: "Identifier", + name: "KEY" + }); + } + } + })] + }).code; + + chai.expect(actualCode).to.be.equal("for (KEY in right);"); + }); + + test("replaceWith (for-loop left expression to variable declaration)", function () { + var expectCode = "for (KEY;;);"; + + var actualCode = transform(expectCode, { + plugins: [new Plugin({ + visitor: { + ForStatement: function (path) { + path.get("init").replaceWith({ + type: "VariableDeclaration", + kind: "var", + declarations: [{ + type: "VariableDeclarator", + id: { + type: "Identifier", + name: "KEY" + } + }] + }); + } + } + })] + }).code; + + chai.expect(actualCode).to.be.equal("for (var KEY;;);"); + }); + + test("replaceWith (for-loop left variable declaration to expression)", function () { + var expectCode = "for (var KEY;;);"; + + var actualCode = transform(expectCode, { + plugins: [new Plugin({ + visitor: { + ForStatement: function (path) { + path.get("init").replaceWith({ + type: "Identifier", + name: "KEY" + }); + } + } + })] + }).code; + + chai.expect(actualCode).to.be.equal("for (KEY;;);"); + }); }); diff --git a/packages/babel-traverse/src/path/introspection.js b/packages/babel-traverse/src/path/introspection.js index 00b27fe2c8..278efc0713 100644 --- a/packages/babel-traverse/src/path/introspection.js +++ b/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 */ diff --git a/packages/babel-traverse/src/path/replacement.js b/packages/babel-traverse/src/path/replacement.js index f1d1555874..8ca15a4c09 100644 --- a/packages/babel-traverse/src/path/replacement.js +++ b/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;