Fix spacing of general (non-specific) tokens + tests

src/client/formatters/lineFormatter.ts

@@ -5,14 +5,15 @@
 import Char from 'typescript-char';
 import { BraceCounter } from '../language/braceCounter';
 import { TextBuilder } from '../language/textBuilder';
+import { TextRangeCollection } from '../language/textRangeCollection';
 import { Tokenizer } from '../language/tokenizer';
 import { ITextRangeCollection, IToken, TokenType } from '../language/types';
 
 export class LineFormatter {
-    private builder: TextBuilder;
-    private tokens: ITextRangeCollection<IToken>;
-    private braceCounter: BraceCounter;
-    private text: string;
+    private builder = new TextBuilder();
+    private tokens: ITextRangeCollection<IToken> = new TextRangeCollection<IToken>([]);
+    private braceCounter = new BraceCounter();
+    private text = '';
 
     // tslint:disable-next-line:cyclomatic-complexity
     public formatLine(text: string): string {
@@ -27,7 +28,7 @@ export class LineFormatter {
 
         const ws = this.text.substr(0, this.tokens.getItemAt(0).start);
         if (ws.length > 0) {
-            this.builder.append(ws); // Preserve leading indentation
+            this.builder.append(ws); // Preserve leading indentation.
         }
 
         for (let i = 0; i < this.tokens.count; i += 1) {
@@ -55,24 +56,28 @@ export class LineFormatter {
                     break;
 
                 case TokenType.Colon:
-                    // x: 1 if not in slice, x[1:y] if inside the slice
+                    // x: 1 if not in slice, x[1:y] if inside the slice.
                     this.builder.append(':');
                     if (!this.braceCounter.isOpened(TokenType.OpenBracket) && (next && next.type !== TokenType.Colon)) {
-                        // Not inside opened [[ ... ] sequence
+                        // Not inside opened [[ ... ] sequence.
                         this.builder.softAppendSpace();
                     }
                     break;
 
                 case TokenType.Comment:
-                    // add space before in-line comment
+                    // Add space before in-line comment.
                     if (prev) {
                         this.builder.softAppendSpace();
                     }
                     this.builder.append(this.text.substring(t.start, t.end));
                     break;
 
+                case TokenType.Semicolon:
+                    this.builder.append(';');
+                    break;
+
                 default:
-                    this.handleOther(t);
+                    this.handleOther(t, i);
                     break;
             }
         }
@@ -85,7 +90,7 @@ export class LineFormatter {
             const opCode = this.text.charCodeAt(t.start);
             switch (opCode) {
                 case Char.Equal:
-                    if (index >= 2 && this.handleEqual(t, index)) {
+                    if (this.handleEqual(t, index)) {
                         return;
                     }
                     break;
@@ -105,27 +110,66 @@ export class LineFormatter {
     }
 
     private handleEqual(t: IToken, index: number): boolean {
-        if (this.braceCounter.isOpened(TokenType.OpenBrace)) {
-            // Check if this is = in function arguments. If so, do not
-            // add spaces around it.
-            const prev = this.tokens.getItemAt(index - 1);
-            const prevPrev = this.tokens.getItemAt(index - 2);
-            if (prev.type === TokenType.Identifier &&
-                (prevPrev.type === TokenType.Comma || prevPrev.type === TokenType.OpenBrace)) {
-                this.builder.append('=');
-                return true;
-            }
+        if (this.isMultipleStatements(index) && !this.braceCounter.isOpened(TokenType.OpenBrace)) {
+            return false; // x = 1; x, y = y, x
+        }
+        // Check if this is = in function arguments. If so, do not add spaces around it.
+        if (this.isEqualsInsideArguments(index)) {
+            this.builder.append('=');
+            return true;
         }
         return false;
     }
 
-    private handleOther(t: IToken): void {
+    private handleOther(t: IToken, index: number): void {
         if (this.isBraceType(t.type)) {
             this.braceCounter.countBrace(t);
+            this.builder.append(this.text.substring(t.start, t.end));
+            return;
         }
+
+        if (this.isEqualsInsideArguments(index - 1)) {
+            // Don't add space around = inside function arguments.
+            this.builder.append(this.text.substring(t.start, t.end));
+            return;
+        }
+
+        if (index > 0) {
+            const prev = this.tokens.getItemAt(index - 1);
+            if (this.isOpenBraceType(prev.type) || prev.type === TokenType.Colon) {
+                // Don't insert space after (, [ or { .
+                this.builder.append(this.text.substring(t.start, t.end));
+                return;
+            }
+        }
+
+        // In general, keep tokens separated.
+        this.builder.softAppendSpace();
         this.builder.append(this.text.substring(t.start, t.end));
     }
 
+    private isEqualsInsideArguments(index: number): boolean {
+        if (index < 1) {
+            return false;
+        }
+        const prev = this.tokens.getItemAt(index - 1);
+        if (prev.type === TokenType.Identifier) {
+            if (index >= 2) {
+                // (x=1 or ,x=1
+                const prevPrev = this.tokens.getItemAt(index - 2);
+                return prevPrev.type === TokenType.Comma || prevPrev.type === TokenType.OpenBrace;
+            } else if (index < this.tokens.count - 2) {
+                const next = this.tokens.getItemAt(index + 1);
+                const nextNext = this.tokens.getItemAt(index + 2);
+                // x=1, or x=1)
+                if (this.isValueType(next.type)) {
+                    return nextNext.type === TokenType.Comma || nextNext.type === TokenType.CloseBrace;
+                }
+            }
+        }
+        return false;
+    }
+
     private isOpenBraceType(type: TokenType): boolean {
         return type === TokenType.OpenBrace || type === TokenType.OpenBracket || type === TokenType.OpenCurly;
     }
@@ -135,4 +179,16 @@ 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) {
+                return true;
+            }
+        }
+        return false;
+    }
 }

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

@@ -65,4 +65,20 @@ suite('Formatting - line formatter', () => {
         const actual = formatter.formatLine('   # comment');
         assert.equal(actual, '   # comment');
     });
+    test('Equals in first argument', () => {
+        const actual = formatter.formatLine('foo(x =0)');
+        assert.equal(actual, 'foo(x=0)');
+    });
+    test('Equals in second argument', () => {
+        const actual = formatter.formatLine('foo(x,y= \"a\",');
+        assert.equal(actual, 'foo(x, y=\"a\",');
+    });
+   test('Equals in multiline arguments', () => {
+        const actual = formatter.formatLine('x = 1,y =-2)');
+        assert.equal(actual, '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)');
+    });
 });

src/test/format/extension.onEnterFormat.test.ts

@@ -59,8 +59,28 @@ suite('Formatting - OnEnter provider', () => {
         assert.equal(text, 'x.y', 'Line ending with period was reformatted');
     });
 
-    test('Formatting line ending in string', async () => {
+    test('Formatting line with unknown neighboring tokens', async () => {
         const text = await formatAtPosition(9, 0);
+        assert.equal(text, 'if x <= 1:', 'Line with unknown neighboring tokens was not formatted');
+    });
+
+    test('Formatting line with unknown neighboring tokens', async () => {
+        const text = await formatAtPosition(10, 0);
+        assert.equal(text, 'if 1 <= x:', 'Line with unknown neighboring tokens was not formatted');
+    });
+
+    test('Formatting method definition with arguments', async () => {
+        const text = await formatAtPosition(11, 0);
+        assert.equal(text, 'def __init__(self, age=23)', 'Method definition with arguments was not formatted');
+    });
+
+    test('Formatting space after open brace', async () => {
+        const text = await formatAtPosition(12, 0);
+        assert.equal(text, 'while(1)', 'Space after open brace was not formatted');
+    });
+
+    test('Formatting line ending in string', async () => {
+        const text = await formatAtPosition(13, 0);
         assert.equal(text, 'x + """', 'Line ending in multiline string was not formatted');
     });
 

src/test/pythonFiles/formatting/fileToFormatOnEnter.py

@@ -6,4 +6,8 @@
 x+1 # 
 @x
 x.y
+if x<=1:
+if 1<=x:
+def __init__(self, age = 23)
+while(1)  
 x+"""