implicit return multiline arrow function inspection (#412)

* add failing test for implicit return arrow fn spanning multiple lines * branching out and retrying a second regex if no bracewrapped fn is found
unexpectedjs · Aug 4, 2017 · 7d84ccd · 7d84ccd
1 parent 938b9da
commit 7d84ccd
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 1 deletion.
diff --git a/lib/types.js b/lib/types.js
@@ -775,6 +775,7 @@ module.exports = function (expect) {
                 body = matchSource[2];
                 var matchBodyAndIndent = body.match(/^(\s*\{)([\s\S]*?)([ ]*)\}\s*$/);
                 var openingBrace;
+                var isWrappedInBraces = true;
                 var closingBrace = '}';
                 if (matchBodyAndIndent) {
                     openingBrace = matchBodyAndIndent[1];
@@ -783,6 +784,17 @@ module.exports = function (expect) {
                     if (bodyIndent.length === 1) {
                         closingBrace = ' }';
                     }
+                } else {
+                    // Attempt to match an arrow function with an implicit return body.
+                    matchBodyAndIndent = body.match(/^(\s*)([\s\S]*?)([ ]*)\s*$/);
+
+                    if (matchBodyAndIndent) {
+                        openingBrace = matchBodyAndIndent[1];
+                        isWrappedInBraces = false;
+                        body = matchBodyAndIndent[2];
+                        bodyIndent = matchBodyAndIndent[3] || '';
+                        closingBrace = '';
+                    }
                 }
 
                 // Remove leading indentation unless the function is a one-liner or it uses multiline string literals
@@ -801,7 +813,7 @@ module.exports = function (expect) {
                     closingBrace = '}';
                 } else if (/^\s*$/.test(body)) {
                     body = '';
-                } else if (/^\s*[^\r\n]{1,30}\s*$/.test(body) && body.indexOf('//') === -1) {
+                } else if (/^\s*[^\r\n]{1,30}\s*$/.test(body) && body.indexOf('//') === -1 && isWrappedInBraces) {
                     body = ' ' + body.trim() + ' ';
                     closingBrace = '}';
                 } else {

diff --git a/test/types/function-type.spec.js b/test/types/function-type.spec.js
@@ -160,6 +160,34 @@ describe('function type', function () {
         });
     }
 
+    // We can't complete this test if the runtime doesn't support arrow functions:
+    var implicitReturnMultilineArrowFunction;
+    try {
+        implicitReturnMultilineArrowFunction = new Function(
+            'return a => \n    a + 1;'
+        )();
+    } catch (e) {}
+
+    if (implicitReturnMultilineArrowFunction) {
+        it('should render an implicit return multiline arrow function', function () {
+            expect(implicitReturnMultilineArrowFunction, 'to inspect as', 'a => \n    a + 1');
+        });
+    }
+
+    // We can't complete this test if the runtime doesn't support arrow functions:
+    var evilImplicitReturnMultilineArrowFunction;
+    try {
+        evilImplicitReturnMultilineArrowFunction = new Function(
+            'return a => \n    a || {};'
+        )();
+    } catch (e) {}
+
+    if (evilImplicitReturnMultilineArrowFunction) {
+        it('should render an implicit return multiline arrow function with an evil alternation', function () {
+            expect(evilImplicitReturnMultilineArrowFunction, 'to inspect as', 'a => \n    a || {}');
+        });
+    }
+
     // We can't complete this test if the runtime doesn't support arrow functions:
     var multiParamArrowFunction;
     try {