#432

Author: DonJayamanne

src/client/typeFormatters/blockFormatProvider.ts

@@ -1,21 +1,65 @@
 import { OnTypeFormattingEditProvider, FormattingOptions, TextEdit, CancellationToken, TextDocument } from 'vscode';
 import { Position } from 'vscode';
-import { CodeBlockFormatProvider } from './contracts';
-import { ElseFormatProvider } from './elseFormatProvider';
+import { CodeBlockFormatProvider } from './codeBlockFormatProvider';
+import { IF_REGEX, ELIF_REGEX, ELSE_REGEX, FOR_IN_REGEX, ASYNC_FOR_IN_REGEX, WHILE_REGEX } from './contracts';
+import { TRY_REGEX, FINALLY_REGEX, EXCEPT_REGEX } from './contracts';
+import { DEF_REGEX, ASYNC_DEF_REGEX, CLASS_REGEX } from './contracts';
 
 export class BlockFormatProviders implements OnTypeFormattingEditProvider {
     private providers: CodeBlockFormatProvider[];
     constructor() {
         this.providers = [];
-        this.providers.push(new ElseFormatProvider());
+        const boundaryBlocks = [
+            DEF_REGEX,
+            ASYNC_DEF_REGEX,
+            CLASS_REGEX
+        ];
+
+        const elseParentBlocks = [
+            IF_REGEX,
+            ELIF_REGEX,
+            FOR_IN_REGEX,
+            ASYNC_FOR_IN_REGEX,
+            WHILE_REGEX,
+            TRY_REGEX,
+            EXCEPT_REGEX
+        ];
+        this.providers.push(new CodeBlockFormatProvider(ELSE_REGEX, elseParentBlocks, boundaryBlocks));
+
+        const elifParentBlocks = [
+            IF_REGEX,
+            ELIF_REGEX
+        ];
+        this.providers.push(new CodeBlockFormatProvider(ELIF_REGEX, elifParentBlocks, boundaryBlocks));
+
+        const exceptParentBlocks = [
+            TRY_REGEX,
+            EXCEPT_REGEX
+        ];
+        this.providers.push(new CodeBlockFormatProvider(EXCEPT_REGEX, exceptParentBlocks, boundaryBlocks));
+
+        const finallyParentBlocks = [
+            TRY_REGEX,
+            EXCEPT_REGEX
+        ];
+        this.providers.push(new CodeBlockFormatProvider(FINALLY_REGEX, finallyParentBlocks, boundaryBlocks));
     }
     provideOnTypeFormattingEdits(document: TextDocument, position: Position, ch: string, options: FormattingOptions, token: CancellationToken): TextEdit[] {
         if (position.line === 0) {
             return [];
         }
 
-        const currentLine = document.lineAt(position.line).text;
-        const provider = this.providers.find(provider => provider.canProvideEdits(currentLine));
+        const currentLine = document.lineAt(position.line);
+        const prevousLine = document.lineAt(position.line - 1);
+
+        // We're only interested in cases where the current block is at the same indentation level as the previous line
+        // E.g. if we have an if..else block, generally the else statement would be at the same level as the code in the if...
+        if (currentLine.firstNonWhitespaceCharacterIndex !== prevousLine.firstNonWhitespaceCharacterIndex) {
+            return [];
+        }
+
+        const currentLineText = currentLine.text;
+        const provider = this.providers.find(provider => provider.canProvideEdits(currentLineText));
         if (provider) {
             return provider.provideEdits(document, position, ch, options, currentLine);
         }

src/client/typeFormatters/codeBlockFormatProvider.ts

@@ -0,0 +1,67 @@
+import { FormattingOptions, TextEdit, TextDocument } from 'vscode';
+import { Position, Range, TextLine } from 'vscode';
+import { BlockRegEx } from './contracts';
+
+export class CodeBlockFormatProvider {
+    constructor(private blockRegExp: BlockRegEx, private previousBlockRegExps: BlockRegEx[], private boundaryRegExps: BlockRegEx[]) {
+    }
+    canProvideEdits(line: string): boolean {
+        return this.blockRegExp.test(line);
+    }
+
+    provideEdits(document: TextDocument, position: Position, ch: string, options: FormattingOptions, line: TextLine): TextEdit[] {
+        // We can have else for the following blocks:
+        // if:
+        // elif x:
+        // for x in y:
+        // while x:
+
+        // We need to find a block statement that is less than or equal to this statement block (but not greater)
+        for (let lineNumber = position.line - 1; lineNumber > 0; lineNumber--) {
+            const prevLine = document.lineAt(lineNumber);
+            const prevLineText = prevLine.text;
+
+            // Oops, we've reached a boundary (like the function or class definition)
+            // Get out of here
+            if (this.boundaryRegExps.some(value => value.test(prevLineText))) {
+                return [];
+            }
+
+            const blockRegEx = this.previousBlockRegExps.find(value => value.test(prevLineText));
+            if (!blockRegEx) {
+                continue;
+            }
+
+            const startOfBlockInLine = prevLine.firstNonWhitespaceCharacterIndex;
+            if (startOfBlockInLine > line.firstNonWhitespaceCharacterIndex) {
+                continue;
+            }
+
+            const startPosition = new Position(position.line, 0);
+            const endPosition = new Position(position.line, line.firstNonWhitespaceCharacterIndex - startOfBlockInLine);
+
+            if (startPosition === endPosition) {
+                return [];
+            }
+            if (options.insertSpaces) {
+                return [
+                    TextEdit.delete(new Range(startPosition, endPosition))
+                ];
+            }
+            else {
+                // Delete everything before the block and insert the same characters we have in the previous block
+                const prefixOfPreviousBlock = prevLineText.substring(0, startOfBlockInLine);
+
+                const startDeletePosition = new Position(position.line, 0);
+                const endDeletePosition = new Position(position.line, line.firstNonWhitespaceCharacterIndex);
+
+                return [
+                    TextEdit.delete(new Range(startDeletePosition, endDeletePosition)),
+                    TextEdit.insert(startDeletePosition, prefixOfPreviousBlock)
+                ];
+            }
+        }
+
+        return [];
+    }
+}

src/client/typeFormatters/contracts.ts

@@ -1,6 +1,3 @@
-import { FormattingOptions, TextEdit, TextDocument } from 'vscode';
-import { Position } from 'vscode';
-
 export class BlockRegEx {
     constructor(private regEx: RegExp, public startWord) {
 
@@ -24,8 +21,3 @@ export const EXCEPT_REGEX = new BlockRegEx(/^( |\t)*except *\w* *(as)? *\w* *: *
 export const DEF_REGEX = new BlockRegEx(/^( |\t)*def \w *\(.*$/g, 'def');
 export const ASYNC_DEF_REGEX = new BlockRegEx(/^( |\t)*async *def \w *\(.*$/g, 'async');
 export const CLASS_REGEX = new BlockRegEx(/^( |\t)*class *\w* *.*: *$/g, 'class');
-
-export interface CodeBlockFormatProvider {
-    canProvideEdits(line: string): boolean;
-    provideEdits(document: TextDocument, position: Position, ch: string, options: FormattingOptions, line: string): TextEdit[];
-}