diff --git a/__snapshots__/test.js.snap b/__snapshots__/test.js.snap index a783540..d33b768 100644 --- a/__snapshots__/test.js.snap +++ b/__snapshots__/test.js.snap @@ -7,7 +7,8 @@ const a = () => { // profile if (true) { - _babel_temp_alias = Math.random(12); + let _babel_temp_alias = Math.random(12); + console.profileEnd(); return _babel_temp_alias; @@ -158,7 +159,9 @@ const b = () => { const a = () => { console.profile(\\"a\\"); - _babel_temp_alias = Math.sqrt(1000000000000); + + let _babel_temp_alias = Math.sqrt(1000000000000); + console.profileEnd(); // profile @@ -166,12 +169,36 @@ const b = () => { }; // profile - _babel_temp_alias = a(); + + let _babel_temp_alias = a(); + console.profileEnd(); return _babel_temp_alias; };" `; +exports[`usage inside classes works inside function properties 1`] = ` +" +class A { + changeVisualState(id) { + console.profile(\\"changeVisualState\\"); + + let _babel_temp_alias = item => { + this.setState({ + newState: {} + }); + callSomeMethod(); + console.profileEnd(); + }; + + console.profileEnd(); + + // profile + return _babel_temp_alias; + } +}" +`; + exports[`wraps a function in profile block 1`] = ` " const b = () => { @@ -184,7 +211,9 @@ const b = () => { const a = () => { console.profile(\\"a\\"); - _babel_temp_alias = b(); + + let _babel_temp_alias = b(); + console.profileEnd(); // profile diff --git a/index.js b/index.js index c884b67..1ddb305 100644 --- a/index.js +++ b/index.js @@ -48,14 +48,27 @@ const generateProfileEnd = () => const generateIdentifier = () => t.identifier('_babel_temp_alias'); -const assignArgument = (identifier, argument) => - t.expressionStatement( - t.assignmentExpression( - "=", - t.identifier('_babel_temp_alias'), - argument - ) - ); +const assignArgument = (identifier, argument, opts) => { + if (opts.blockDepth === 1) { + return t.variableDeclaration( + 'let', + [ + t.variableDeclarator( + t.identifier('_babel_temp_alias'), + argument + ) + ] + ); + } else { + return t.expressionStatement( + t.assignmentExpression( + "=", + t.identifier('_babel_temp_alias'), + argument + ) + ); + } +} const SkipVisitor = function (path, args) { @@ -68,6 +81,7 @@ const ReturnVisitor = { const argument = path.node.argument; args.state.gotReturn = true; + args.state.blockDepth++; if (args.state.gotProfileComment) { if (!argument || argument.type === 'Literal' || argument.type === 'NumericLiteral' || argument.type === 'StringLiteral' || argument.type === 'Identifier') { @@ -81,7 +95,9 @@ const ReturnVisitor = { // set the argument of return to the variable identifier const identifier = generateIdentifier(); - const assignment = assignArgument(identifier, argument); + const assignment = assignArgument(identifier, argument, { + blockDepth: args.state.blockDepth + }); path.insertBefore(assignment); path.insertBefore(generateProfileEnd()); path.node.argument = identifier; @@ -139,7 +155,8 @@ const ConsoleProfileVisitor = function (babel) { BlockStatement: function (path, args) { let state = { gotProfileComment: false, - gotReturn: false + gotReturn: false, + blockDepth: 0 }; BlockVisitor(path, { state }); diff --git a/test.js b/test.js index 0851c0b..ead9fa8 100644 --- a/test.js +++ b/test.js @@ -170,3 +170,25 @@ it('maintains indentation when for if blocks', () => { expect(code).toMatchSnapshot(); }); + +describe('usage inside classes', () => { + it('works inside function properties', () => { + const example = ` + class A { + changeVisualState (id) { + // profile + return item => { + this.setState({ + newState: {} + }); + callSomeMethod(); + }; + } + } + `; + + const code = babel.transform(example, { plugins: [ plugin ] }).code; + expect(code).toMatchSnapshot(); + }); +}); +