Fixed generator shorthand in class block

Fixes #1013
beautifier · Dec 22, 2016 · e8f01d6 · e8f01d6
1 parent 387708b
commit e8f01d6
Show file tree

Hide file tree

Showing 5 changed files with 81 additions and 5 deletions.
diff --git a/js/lib/beautify.js b/js/lib/beautify.js
@@ -1312,7 +1312,9 @@ if (!Object.values) {
                 var in_ternary = false;
                 var isGeneratorAsterisk = current_token.text === '*' &&
                     ((last_type === 'TK_RESERVED' && in_array(flags.last_text, ['function', 'yield'])) ||
-                        (flags.mode === MODE.ObjectLiteral && in_array(last_type, ['TK_START_BLOCK', 'TK_COMMA'])));
+                        (flags.mode === MODE.ObjectLiteral && in_array(last_type, ['TK_START_BLOCK', 'TK_COMMA'])) ||
+                        (flags.mode === MODE.BlockStatement && in_array(last_type, ['TK_START_BLOCK', 'TK_COMMA', 'TK_END_BLOCK', 'TK_SEMICOLON']))
+                    );
                 var isUnary = in_array(current_token.text, ['-', '+']) && (
                     in_array(last_type, ['TK_START_BLOCK', 'TK_START_EXPR', 'TK_EQUALS', 'TK_OPERATOR']) ||
                     in_array(flags.last_text, Tokenizer.line_starters) ||
@@ -1391,7 +1393,8 @@ if (!Object.values) {
                 if (isGeneratorAsterisk) {
                     allow_wrap_or_preserved_newline();
                     space_before = false;
-                    space_after = false;
+                    var next_token = get_token(1);
+                    space_after = next_token && in_array(next_token.type, ['TK_WORD', 'TK_RESERVED']);
                 } else if (current_token.text === '...') {
                     allow_wrap_or_preserved_newline();
                     space_before = last_type === 'TK_START_BLOCK';

diff --git a/js/test/generated/beautify-javascript-tests.js b/js/test/generated/beautify-javascript-tests.js
@@ -379,6 +379,28 @@ function run_javascript_tests(test_obj, Urlencoded, js_beautify, html_beautify,
             '        yield 42;\n' +
             '    }\n' +
             '};');
+
+        // also handle generator shorthand in class - #1013
+        bt(
+            'class A {\n' +
+            '    fn() {\n' +
+            '        return true;\n' +
+            '    }\n' +
+            '\n' +
+            '    * gen() {\n' +
+            '        return true;\n' +
+            '    }\n' +
+            '}');
+        bt(
+            'class A {\n' +
+            '    * gen() {\n' +
+            '        return true;\n' +
+            '    }\n' +
+            '\n' +
+            '    fn() {\n' +
+            '        return true;\n' +
+            '    }\n' +
+            '}');
 
 
         //============================================================

diff --git a/python/jsbeautifier/__init__.py b/python/jsbeautifier/__init__.py
@@ -1232,7 +1232,9 @@ def handle_operator(self, current_token):
         in_ternary = False
         isGeneratorAsterisk = current_token.text == '*' and \
             ((self.last_type == 'TK_RESERVED' and self.flags.last_text in ['function', 'yield']) or
-                (self.flags.mode == MODE.ObjectLiteral and self.last_type in ['TK_START_BLOCK', 'TK_COMMA']))
+                (self.flags.mode == MODE.ObjectLiteral and self.last_type in ['TK_START_BLOCK', 'TK_COMMA']) or
+                (self.flags.mode == MODE.BlockStatement and self.last_type in ['TK_START_BLOCK', 'TK_COMMA', 'TK_END_BLOCK', 'TK_SEMICOLON'])
+                )
         isUnary = current_token.text in ['+', '-'] \
             and (self.last_type in ['TK_START_BLOCK', 'TK_START_EXPR', 'TK_EQUALS', 'TK_OPERATOR'] \
             or self.flags.last_text in Tokenizer.line_starters or self.flags.last_text == ',')
@@ -1303,7 +1305,8 @@ def handle_operator(self, current_token):
         if isGeneratorAsterisk:
             self.allow_wrap_or_preserved_newline(current_token)
             space_before = False
-            space_after = False
+            next_token = self.get_token(1)
+            space_after = next_token and next_token.type in ['TK_WORD','TK_RESERVED']
         elif current_token.text == '...':
             self.allow_wrap_or_preserved_newline(current_token)
             space_before = self.last_type == 'TK_START_BLOCK'

diff --git a/python/jsbeautifier/tests/generated/tests.py b/python/jsbeautifier/tests/generated/tests.py
@@ -207,6 +207,28 @@ def unicode_char(value):
             '        yield 42;\n' +
             '    }\n' +
             '};')
+
+        # also handle generator shorthand in class - #1013
+        bt(
+            'class A {\n' +
+            '    fn() {\n' +
+            '        return true;\n' +
+            '    }\n' +
+            '\n' +
+            '    * gen() {\n' +
+            '        return true;\n' +
+            '    }\n' +
+            '}')
+        bt(
+            'class A {\n' +
+            '    * gen() {\n' +
+            '        return true;\n' +
+            '    }\n' +
+            '\n' +
+            '    fn() {\n' +
+            '        return true;\n' +
+            '    }\n' +
+            '}')
 
 
         #============================================================

diff --git a/test/data/javascript/tests.js b/test/data/javascript/tests.js
@@ -119,7 +119,8 @@ exports.test_data = {
         tests: [
             { unchanged: 'return {\n    foo() {\n        return 42;\n    }\n}' },
             {
-                unchanged: ['var foo = {',
+                unchanged: [
+                    'var foo = {',
                     '    * bar() {',
                     '        yield 42;',
                     '    }',
@@ -137,6 +138,31 @@ exports.test_data = {
                     '    }',
                     '};'
                 ]
+            }, {
+                comment: 'also handle generator shorthand in class - #1013',
+                unchanged: [
+                    'class A {',
+                    '    fn() {',
+                    '        return true;',
+                    '    }',
+                    '',
+                    '    * gen() {',
+                    '        return true;',
+                    '    }',
+                    '}'
+                ]
+            }, {
+                unchanged: [
+                    'class A {',
+                    '    * gen() {',
+                    '        return true;',
+                    '    }',
+                    '',
+                    '    fn() {',
+                    '        return true;',
+                    '    }',
+                    '}'
+                ]
             }
         ]
     }, {