From 442cc05cec89e222434fae8e97d2f78c67dbd591 Mon Sep 17 00:00:00 2001
From: Liam Newman <bitwiseman@gmail.com>
Date: Tue, 26 Jan 2016 16:35:44 -0800
Subject: [PATCH] Arrow functions can have StatementBlocks

Fixes #854
---
 js/lib/beautify.js                             |  3 +++
 js/test/generated/beautify-javascript-tests.js | 12 ++++++++++++
 python/js-beautify-test.py                     |  8 ++++++--
 python/jsbeautifier/__init__.py                |  7 +++++--
 python/jsbeautifier/tests/generated/tests.py   | 12 ++++++++++++
 test/data/javascript/tests.js                  | 14 ++++++++++++++
 6 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/js/lib/beautify.js b/js/lib/beautify.js
index a5358744..4de4092b 100644
--- a/js/lib/beautify.js
+++ b/js/lib/beautify.js
@@ -766,6 +766,9 @@
                 } else {
                     set_mode(MODE.BlockStatement);
                 }
+            } else if (last_type === 'TK_OPERATOR' && flags.last_text === '=>') {
+                // arrow function: (param1, paramN) => { statements }
+                set_mode(MODE.BlockStatement);
             } else if (in_array(last_type, ['TK_EQUALS', 'TK_START_EXPR', 'TK_COMMA', 'TK_OPERATOR']) ||
                 (last_type === 'TK_RESERVED' && in_array(flags.last_text, ['return', 'throw', 'import']))
                 ) {
diff --git a/js/test/generated/beautify-javascript-tests.js b/js/test/generated/beautify-javascript-tests.js
index d2310fdb..3fb6fa4d 100644
--- a/js/test/generated/beautify-javascript-tests.js
+++ b/js/test/generated/beautify-javascript-tests.js
@@ -1416,6 +1416,18 @@ function run_javascript_tests(test_obj, Urlencoded, js_beautify, html_beautify,
             '        return 1;\n' +
             '    }\n' +
             '}.fn2()');
+        
+        // Issue 854 - Arrow function with statement block
+        bt(
+            'test(() => {\n' +
+            '    var a = {}\n' +
+            '\n' +
+            '    a.what = () => true ? 1 : 2\n' +
+            '\n' +
+            '    a.thing = () => {\n' +
+            '        b();\n' +
+            '    }\n' +
+            '})');
 
 
 
diff --git a/python/js-beautify-test.py b/python/js-beautify-test.py
index bd373bce..3c215b12 100755
--- a/python/js-beautify-test.py
+++ b/python/js-beautify-test.py
@@ -1,7 +1,11 @@
 #!/usr/bin/env python
 
+import sys
 import unittest
 
-if __name__ == "__main__":
+def run_tests():
     suite = unittest.TestLoader().discover('.', pattern = "test*.py")
-    unittest.TextTestRunner(verbosity=2).run(suite)
+    return unittest.TextTestRunner(verbosity=2).run(suite)
+
+if __name__ == "__main__":
+    sys.exit(not run_tests().wasSuccessful())
diff --git a/python/jsbeautifier/__init__.py b/python/jsbeautifier/__init__.py
index e72d23c8..d42aaaa3 100644
--- a/python/jsbeautifier/__init__.py
+++ b/python/jsbeautifier/__init__.py
@@ -736,16 +736,19 @@ class Beautifier:
             # We don't support TypeScript,but we didn't break it for a very long time.
             # We'll try to keep not breaking it.
             if not self.last_last_text in ['class','interface']:
-                self.set_mode(MODE.ObjectLiteral);
+                self.set_mode(MODE.ObjectLiteral)
             else:
                 self.set_mode(MODE.BlockStatement)
+        elif self.last_type == 'TK_OPERATOR' and self.flags.last_text == '=>':
+            # arrow function: (param1, paramN) => { statements }
+            self.set_mode(MODE.BlockStatement)
         elif self.last_type in ['TK_EQUALS', 'TK_START_EXPR', 'TK_COMMA', 'TK_OPERATOR'] or \
             (self.last_type == 'TK_RESERVED' and self.flags.last_text in ['return', 'throw', 'import']):
             # Detecting shorthand function syntax is difficult by scanning forward,
             #     so check the surrounding context.
             # If the block is being returned, imported, passed as arg,
             #     assigned with = or assigned in a nested object, treat as an ObjectLiteral.
-            self.set_mode(MODE.ObjectLiteral);
+            self.set_mode(MODE.ObjectLiteral)
         else:
             self.set_mode(MODE.BlockStatement)
 
diff --git a/python/jsbeautifier/tests/generated/tests.py b/python/jsbeautifier/tests/generated/tests.py
index 042dda81..db2f2d18 100644
--- a/python/jsbeautifier/tests/generated/tests.py
+++ b/python/jsbeautifier/tests/generated/tests.py
@@ -1191,6 +1191,18 @@ class TestJSBeautifier(unittest.TestCase):
             '        return 1;\n' +
             '    }\n' +
             '}.fn2()')
+        
+        # Issue 854 - Arrow function with statement block
+        bt(
+            'test(() => {\n' +
+            '    var a = {}\n' +
+            '\n' +
+            '    a.what = () => true ? 1 : 2\n' +
+            '\n' +
+            '    a.thing = () => {\n' +
+            '        b();\n' +
+            '    }\n' +
+            '})')
 
         # Destructured and related
         self.options.brace_style = 'collapse-preserve-inline'
diff --git a/test/data/javascript/tests.js b/test/data/javascript/tests.js
index 47d1f6f2..849f3e49 100644
--- a/test/data/javascript/tests.js
+++ b/test/data/javascript/tests.js
@@ -1420,6 +1420,20 @@ exports.test_data = {
                     '    }',
                     '}.fn2()'
                 ]
+            },
+            {
+                comment: "Issue 854 - Arrow function with statement block",
+                unchanged: [
+                    'test(() => {',
+                    '    var a = {}',
+                    '',
+                    '    a.what = () => true ? 1 : 2',
+                    '',
+                    '    a.thing = () => {',
+                    '        b();',
+                    '    }',
+                    '})'
+                ]
             }
         ]
     }, {