Improve detection of the function argument list in autoformat

src/client/formatters/lineFormatter.ts

@@ -3,6 +3,7 @@
 
 // tslint:disable-next-line:import-name
 import Char from 'typescript-char';
+import { TextDocument } from 'vscode';
 import { BraceCounter } from '../language/braceCounter';
 import { TextBuilder } from '../language/textBuilder';
 import { TextRangeCollection } from '../language/textRangeCollection';
@@ -14,11 +15,15 @@ export class LineFormatter {
     private tokens: ITextRangeCollection<IToken> = new TextRangeCollection<IToken>([]);
     private braceCounter = new BraceCounter();
     private text = '';
+    private document?: TextDocument;
+    private lineNumber = 0;
 
     // tslint:disable-next-line:cyclomatic-complexity
-    public formatLine(text: string): string {
-        this.tokens = new Tokenizer().tokenize(text);
-        this.text = text;
+    public formatLine(document: TextDocument, lineNumber: number): string {
+        this.document = document;
+        this.lineNumber = lineNumber;
+        this.text = document.lineAt(lineNumber).text;
+        this.tokens = new Tokenizer().tokenize(this.text);
         this.builder = new TextBuilder();
         this.braceCounter = new BraceCounter();
 
@@ -107,7 +112,7 @@ export class LineFormatter {
                     this.builder.append(this.text[t.start]);
                     return;
                 case Char.Asterisk:
-                    if (prev && prev.type === TokenType.Identifier && prev.length === 6 && this.text.substr(prev.start, prev.length) === 'lambda') {
+                    if (prev && this.isKeyword(prev, 'lambda')) {
                         this.builder.softAppendSpace();
                         this.builder.append('*');
                         return;
@@ -122,7 +127,7 @@ export class LineFormatter {
                     this.builder.append('**');
                     return;
                 }
-                if (prev && prev.type === TokenType.Identifier && prev.length === 6 && this.text.substr(prev.start, prev.length) === 'lambda') {
+                if (prev && this.isKeyword(prev, 'lambda')) {
                     this.builder.softAppendSpace();
                     this.builder.append('**');
                     return;
@@ -194,6 +199,8 @@ export class LineFormatter {
             this.builder.softAppendSpace();
         }
     }
+
+    // tslint:disable-next-line:cyclomatic-complexity
     private isEqualsInsideArguments(index: number): boolean {
         // Since we don't have complete statement, this is mostly heuristics.
         // Therefore the code may not be handling all possible ways of the
@@ -217,28 +224,31 @@ export class LineFormatter {
             return true; // Line ends in comma
         }
 
-        if (index >= 2) {
-            // (x=1 or ,x=1
-            const prevPrev = this.tokens.getItemAt(index - 2);
-            return prevPrev.type === TokenType.Comma || prevPrev.type === TokenType.OpenBrace;
+        if (last.type === TokenType.Comment && this.tokens.count > 1 && this.tokens.getItemAt(this.tokens.count - 2).type === TokenType.Comma) {
+            return true; // Line ends in comma and then comment
         }
 
-        if (index >= this.tokens.count - 2) {
-            return false;
+        if (this.document) {
+            const prevLine = this.lineNumber > 0 ? this.document.lineAt(this.lineNumber - 1).text : '';
+            const prevLineTokens = new Tokenizer().tokenize(prevLine);
+            if (prevLineTokens.count > 0) {
+                const lastOnPrevLine = prevLineTokens.getItemAt(prevLineTokens.count - 1);
+                if (lastOnPrevLine.type === TokenType.Comma) {
+                    return true; // Previous line ends in comma
+                }
+                if (lastOnPrevLine.type === TokenType.Comment && prevLineTokens.count > 1 && prevLineTokens.getItemAt(prevLineTokens.count - 2).type === TokenType.Comma) {
+                    return true; // Previous line ends in comma and then comment
+                }
+            }
         }
 
-        const next = this.tokens.getItemAt(index + 1);
-        const nextNext = this.tokens.getItemAt(index + 2);
-        // x=1, or x=1)
-        if (this.isValueType(next.type)) {
-            if (nextNext.type === TokenType.CloseBrace) {
+        for (let i = 0; i < index; i += 1) {
+            const t = this.tokens.getItemAt(i);
+            if (this.isKeyword(t, 'lambda')) {
                 return true;
             }
-            if (nextNext.type === TokenType.Comma) {
-                return last.type === TokenType.CloseBrace;
-            }
         }
-        return false;
+        return this.braceCounter.isOpened(TokenType.OpenBrace);
     }
 
     private isOpenBraceType(type: TokenType): boolean {
@@ -250,10 +260,6 @@ export class LineFormatter {
     private isBraceType(type: TokenType): boolean {
         return this.isOpenBraceType(type) || this.isCloseBraceType(type);
     }
-    private isValueType(type: TokenType): boolean {
-        return type === TokenType.Identifier || type === TokenType.Unknown ||
-            type === TokenType.Number || type === TokenType.String;
-    }
     private isMultipleStatements(index: number): boolean {
         for (let i = index; i >= 0; i -= 1) {
             if (this.tokens.getItemAt(i).type === TokenType.Semicolon) {
@@ -268,4 +274,7 @@ export class LineFormatter {
             s === 'import' || s === 'except' || s === 'for' ||
             s === 'as' || s === 'is';
     }
+    private isKeyword(t: IToken, keyword: string): boolean {
+        return t.type === TokenType.Identifier && t.length === keyword.length && this.text.substr(t.start, t.length) === keyword;
+    }
 }

src/client/typeFormatters/onEnterFormatter.ts

@@ -1,20 +1,20 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import * as vscode from 'vscode';
+import { CancellationToken, FormattingOptions, OnTypeFormattingEditProvider, Position, TextDocument, TextEdit } from 'vscode';
 import { LineFormatter } from '../formatters/lineFormatter';
 import { TokenizerMode, TokenType } from '../language/types';
 import { getDocumentTokens } from '../providers/providerUtilities';
 
-export class OnEnterFormatter implements vscode.OnTypeFormattingEditProvider {
+export class OnEnterFormatter implements OnTypeFormattingEditProvider {
     private readonly formatter = new LineFormatter();
 
     public provideOnTypeFormattingEdits(
-        document: vscode.TextDocument,
-        position: vscode.Position,
+        document: TextDocument,
+        position: Position,
         ch: string,
-        options: vscode.FormattingOptions,
-        cancellationToken: vscode.CancellationToken): vscode.TextEdit[] {
+        options: FormattingOptions,
+        cancellationToken: CancellationToken): TextEdit[] {
         if (position.line === 0) {
             return [];
         }
@@ -30,10 +30,10 @@ export class OnEnterFormatter implements vscode.OnTypeFormattingEditProvider {
                 return [];
             }
         }
-        const formatted = this.formatter.formatLine(prevLine.text);
+        const formatted = this.formatter.formatLine(document, prevLine.lineNumber);
         if (formatted === prevLine.text) {
             return [];
         }
-        return [new vscode.TextEdit(prevLine.range, formatted)];
+        return [new TextEdit(prevLine.range, formatted)];
     }
 }

src/test/format/extension.lineFormatter.test.ts

@@ -5,6 +5,8 @@
 import * as assert from 'assert';
 import * as fs from 'fs';
 import * as path from 'path';
+import * as TypeMoq from 'typemoq';
+import { TextDocument, TextLine } from 'vscode';
 import '../../client/common/extensions';
 import { LineFormatter } from '../../client/formatters/lineFormatter';
 
@@ -17,124 +19,142 @@ suite('Formatting - line formatter', () => {
     const formatter = new LineFormatter();
 
     test('Operator spacing', () => {
-        const actual = formatter.formatLine('( x  +1 )*y/ 3');
-        assert.equal(actual, '(x + 1) * y / 3');
+        testFormatLine('( x  +1 )*y/ 3', '(x + 1) * y / 3');
     });
     test('Braces spacing', () => {
-        const actual = formatter.formatLine('foo =(0 ,)');
-        assert.equal(actual, 'foo = (0,)');
+        testFormatLine('foo =(0 ,)', 'foo = (0,)');
     });
     test('Function arguments', () => {
-        const actual = formatter.formatLine('foo (0 , x= 1, (3+7) , y , z )');
-        assert.equal(actual, 'foo(0, x=1, (3 + 7), y, z)');
+        testFormatLine('z=foo (0 , x= 1, (3+7) , y , z )',
+            'z = foo(0, x=1, (3 + 7), y, z)');
     });
     test('Colon regular', () => {
-        const actual = formatter.formatLine('if x == 4 : print x,y; x,y= y, x');
-        assert.equal(actual, 'if x == 4: print x, y; x, y = y, x');
+        testFormatLine('if x == 4 : print x,y; x,y= y, x',
+            'if x == 4: print x, y; x, y = y, x');
     });
     test('Colon slices', () => {
-        const actual = formatter.formatLine('x[1: 30]');
-        assert.equal(actual, 'x[1:30]');
+        testFormatLine('x[1: 30]', 'x[1:30]');
     });
     test('Colon slices in arguments', () => {
-        const actual = formatter.formatLine('spam ( ham[ 1 :3], {eggs : 2})');
-        assert.equal(actual, 'spam(ham[1:3], {eggs: 2})');
+        testFormatLine('spam ( ham[ 1 :3], {eggs : 2})',
+            'spam(ham[1:3], {eggs: 2})');
     });
     test('Colon slices with double colon', () => {
-        const actual = formatter.formatLine('ham [1:9 ], ham[ 1: 9:   3], ham[: 9 :3], ham[1: :3], ham [ 1: 9:]');
-        assert.equal(actual, 'ham[1:9], ham[1:9:3], ham[:9:3], ham[1::3], ham[1:9:]');
+        testFormatLine('ham [1:9 ], ham[ 1: 9:   3], ham[: 9 :3], ham[1: :3], ham [ 1: 9:]',
+            'ham[1:9], ham[1:9:3], ham[:9:3], ham[1::3], ham[1:9:]');
     });
     test('Colon slices with operators', () => {
-        const actual = formatter.formatLine('ham [lower+ offset :upper+offset]');
-        assert.equal(actual, 'ham[lower + offset:upper + offset]');
+        testFormatLine('ham [lower+ offset :upper+offset]',
+            'ham[lower + offset:upper + offset]');
     });
     test('Colon slices with functions', () => {
-        const actual = formatter.formatLine('ham[ : upper_fn ( x) : step_fn(x )], ham[ :: step_fn(x)]');
-        assert.equal(actual, 'ham[:upper_fn(x):step_fn(x)], ham[::step_fn(x)]');
+        testFormatLine('ham[ : upper_fn ( x) : step_fn(x )], ham[ :: step_fn(x)]',
+            'ham[:upper_fn(x):step_fn(x)], ham[::step_fn(x)]');
     });
     test('Colon in for loop', () => {
-        const actual = formatter.formatLine('for index in  range( len(fruits) ): ');
-        assert.equal(actual, 'for index in range(len(fruits)):');
+        testFormatLine('for index in  range( len(fruits) ): ',
+            'for index in range(len(fruits)):');
     });
     test('Nested braces', () => {
-        const actual = formatter.formatLine('[ 1 :[2: (x,),y]]{1}');
-        assert.equal(actual, '[1:[2:(x,), y]]{1}');
+        testFormatLine('[ 1 :[2: (x,),y]]{1}', '[1:[2:(x,), y]]{1}');
     });
     test('Trailing comment', () => {
-        const actual = formatter.formatLine('x=1  # comment');
-        assert.equal(actual, 'x = 1 # comment');
+        testFormatLine('x=1  # comment', 'x = 1 # comment');
     });
     test('Single comment', () => {
-        const actual = formatter.formatLine('# comment');
-        assert.equal(actual, '# comment');
+        testFormatLine('# comment', '# comment');
     });
     test('Comment with leading whitespace', () => {
-        const actual = formatter.formatLine('   # comment');
-        assert.equal(actual, '   # comment');
+        testFormatLine('   # comment', '   # comment');
     });
     test('Equals in first argument', () => {
-        const actual = formatter.formatLine('foo(x =0)');
-        assert.equal(actual, 'foo(x=0)');
+        testFormatLine('foo(x =0)', 'foo(x=0)');
     });
     test('Equals in second argument', () => {
-        const actual = formatter.formatLine('foo(x,y= \"a\",');
-        assert.equal(actual, 'foo(x, y=\"a\",');
+        testFormatLine('foo(x,y= \"a\",', 'foo(x, y=\"a\",');
     });
     test('Equals in multiline arguments', () => {
-        const actual = formatter.formatLine('x = 1,y =-2)');
-        assert.equal(actual, 'x=1, y=-2)');
+        testFormatLine2('foo(a,', 'x = 1,y =-2)', 'x=1, y=-2)');
     });
     test('Equals in multiline arguments starting comma', () => {
-        const actual = formatter.formatLine(',x = 1,y =m)');
-        assert.equal(actual, ', x=1, y=m)');
+        testFormatLine(',x = 1,y =m)', ', x=1, y=m)');
     });
     test('Equals in multiline arguments ending comma', () => {
-        const actual = formatter.formatLine('x = 1,y =m,');
-        assert.equal(actual, 'x=1, y=m,');
+        testFormatLine('x = 1,y =m,', 'x=1, y=m,');
     });
     test('Operators without following space', () => {
-        const actual = formatter.formatLine('foo( *a, ** b, ! c)');
-        assert.equal(actual, 'foo(*a, **b, !c)');
+        testFormatLine('foo( *a, ** b, ! c)', 'foo(*a, **b, !c)');
     });
     test('Brace after keyword', () => {
-        const actual = formatter.formatLine('for x in(1,2,3)');
-        assert.equal(actual, 'for x in (1, 2, 3)');
+        testFormatLine('for x in(1,2,3)', 'for x in (1, 2, 3)');
     });
     test('Dot operator', () => {
-        const actual = formatter.formatLine('x.y');
-        assert.equal(actual, 'x.y');
+        testFormatLine('x.y', 'x.y');
     });
     test('Unknown tokens no space', () => {
-        const actual = formatter.formatLine('abc\\n\\');
-        assert.equal(actual, 'abc\\n\\');
+        testFormatLine('abc\\n\\', 'abc\\n\\');
     });
     test('Unknown tokens with space', () => {
-        const actual = formatter.formatLine('abc \\n \\');
-        assert.equal(actual, 'abc \\n \\');
+        testFormatLine('abc \\n \\', 'abc \\n \\');
     });
     test('Double asterisk', () => {
-        const actual = formatter.formatLine('a**2, ** k');
-        assert.equal(actual, 'a ** 2, **k');
+        testFormatLine('a**2, ** k', 'a ** 2, **k');
     });
     test('Lambda', () => {
-        const actual = formatter.formatLine('lambda * args, :0');
-        assert.equal(actual, 'lambda *args,: 0');
+        testFormatLine('lambda * args, :0', 'lambda *args,: 0');
     });
     test('Comma expression', () => {
-        const actual = formatter.formatLine('x=1,2,3');
-        assert.equal(actual, 'x = 1, 2, 3');
+        testFormatLine('x=1,2,3', 'x = 1, 2, 3');
     });
     test('is exression', () => {
-        const actual = formatter.formatLine('a( (False is  2)  is 3)');
-        assert.equal(actual, 'a((False is 2) is 3)');
+        testFormatLine('a( (False is  2)  is 3)', 'a((False is 2) is 3)');
+    });
+    test('Function returning tuple', () => {
+        testFormatLine('x,y=f(a)', 'x, y = f(a)');
     });
     test('Grammar file', () => {
         const content = fs.readFileSync(grammarFile).toString('utf8');
         const lines = content.splitLines({ trim: false, removeEmptyEntries: false });
+        let prevLine = '';
         for (let i = 0; i < lines.length; i += 1) {
             const line = lines[i];
-            const actual = formatter.formatLine(line);
-            assert.equal(actual, line, `Line ${i + 1} changed: '${line}' to '${actual}'`);
+            const actual = formatLine2(prevLine, line);
+            assert.equal(actual, line, `Line ${i + 1} changed: '${line.trim()}' to '${actual.trim()}'`);
+            prevLine = line;
         }
     });
+
+    function testFormatLine(text: string, expected: string): void {
+        const actual = formatLine(text);
+        assert.equal(actual, expected);
+    }
+
+    function formatLine(text: string): string {
+        const line = TypeMoq.Mock.ofType<TextLine>();
+        line.setup(x => x.text).returns(() => text);
+
+        const document = TypeMoq.Mock.ofType<TextDocument>();
+        document.setup(x => x.lineAt(TypeMoq.It.isAnyNumber())).returns(() => line.object);
+
+        return formatter.formatLine(document.object, 0);
+    }
+
+    function formatLine2(prevLineText: string, lineText: string): string {
+        const thisLine = TypeMoq.Mock.ofType<TextLine>();
+        thisLine.setup(x => x.text).returns(() => lineText);
+
+        const prevLine = TypeMoq.Mock.ofType<TextLine>();
+        prevLine.setup(x => x.text).returns(() => prevLineText);
+
+        const document = TypeMoq.Mock.ofType<TextDocument>();
+        document.setup(x => x.lineAt(0)).returns(() => prevLine.object);
+        document.setup(x => x.lineAt(1)).returns(() => thisLine.object);
+
+        return formatter.formatLine(document.object, 1);
+    }
+
+    function testFormatLine2(prevLineText: string, lineText: string, expected: string): void {
+        const actual = formatLine2(prevLineText, lineText);
+        assert.equal(actual, expected);
+    }
 });

src/test/pythonFiles/formatting/pythonGrammar.py

@@ -646,7 +646,7 @@ def test_lambdef(self):
         l2 = lambda: a[d] # XXX just testing the expression
         l3 = lambda: [2 < x for x in [-1, 3, 0]]
         self.assertEqual(l3(), [0, 1, 0])
-        l4 = lambda x = lambda y = lambda z = 1: z: y(): x()
+        l4 = lambda x=lambda y=lambda z=1: z: y(): x()
         self.assertEqual(l4(), 1)
         l5 = lambda x, y, z=2: x + y + z
         self.assertEqual(l5(1, 2), 5)
@@ -696,8 +696,8 @@ def test_expr_stmt(self):
         x = 1
         x = 1, 2, 3
         x = y = z = 1, 2, 3
-        x, y, z=1, 2, 3
-        abc=a, b, c=x, y, z=xyz = 1, 2, (3, 4)
+        x, y, z = 1, 2, 3
+        abc = a, b, c = x, y, z = xyz = 1, 2, (3, 4)
 
         check_syntax_error(self, "x + 1 = 1")
         check_syntax_error(self, "a + 1 = b + 2")
@@ -730,7 +730,7 @@ def test_former_statements_refer_to_builtins(self):
     def test_del_stmt(self):
         # 'del' exprlist
         abc = [1, 2, 3]
-        x, y, z=abc
+        x, y, z = abc
         xyz = x, y, z
 
         del abc