#432

Author: DonJayamanne

src/client/extension.ts

@@ -22,6 +22,7 @@ import * as jup from './jupyter/main';
 import { HelpProvider } from './helpProvider';
 import { activateFormatOnSaveProvider } from './providers/formatOnSaveProvider';
 import { WorkspaceSymbols } from './workspaceSymbols/main';
+import { BlockFormatProviders } from './typeFormatters/blockFormatProvider';
 import * as os from 'os';
 
 const PYTHON: vscode.DocumentFilter = { language: 'python', scheme: 'file' };
@@ -95,6 +96,7 @@ export function activate(context: vscode.ExtensionContext) {
 
     context.subscriptions.push(new WorkspaceSymbols(lintingOutChannel));
 
+    context.subscriptions.push(vscode.languages.registerOnTypeFormattingEditProvider(PYTHON, new BlockFormatProviders(), ':'));
     // In case we have CR LF
     const triggerCharacters: string[] = os.EOL.split('');
     triggerCharacters.shift();

src/client/typeFormatters/blockFormatProvider.ts

@@ -0,0 +1,25 @@
+import { OnTypeFormattingEditProvider, FormattingOptions, TextEdit, CancellationToken, TextDocument } from 'vscode';
+import { Position } from 'vscode';
+import { CodeBlockFormatProvider } from './contracts';
+import { ElseFormatProvider } from './elseFormatProvider';
+
+export class BlockFormatProviders implements OnTypeFormattingEditProvider {
+    private providers: CodeBlockFormatProvider[];
+    constructor() {
+        this.providers = [];
+        this.providers.push(new ElseFormatProvider());
+    }
+    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));
+        if (provider) {
+            return provider.provideEdits(document, position, ch, options, currentLine);
+        }
+
+        return [];
+    }
+}

src/client/typeFormatters/contracts.ts

@@ -0,0 +1,31 @@
+import { FormattingOptions, TextEdit, TextDocument } from 'vscode';
+import { Position } from 'vscode';
+
+export class BlockRegEx {
+    constructor(private regEx: RegExp, public startWord) {
+
+    }
+    public test(value: string): boolean {
+        // Clear the cache
+        this.regEx.lastIndex = -1;
+        return this.regEx.test(value);
+    }
+}
+
+export const IF_REGEX = new BlockRegEx(/^( |\t)*if +.*: *$/g, 'if');
+export const ELIF_REGEX = new BlockRegEx(/^( |\t)*elif +.*: *$/g, 'elif');
+export const ELSE_REGEX = new BlockRegEx(/^( |\t)*else *: *$/g, 'else');
+export const FOR_IN_REGEX = new BlockRegEx(/^( |\t)*for \w in .*: *$/g, 'for');
+export const ASYNC_FOR_IN_REGEX = new BlockRegEx(/^( |\t)*async *for \w in .*: *$/g, 'for');
+export const WHILE_REGEX = new BlockRegEx(/^( |\t)*while .*: *$/g, 'while');
+export const TRY_REGEX = new BlockRegEx(/^( |\t)*try *: *$/g, 'try');
+export const FINALLY_REGEX = new BlockRegEx(/^( |\t)*finally *: *$/g, 'finally');
+export const EXCEPT_REGEX = new BlockRegEx(/^( |\t)*except *\w* *(as)? *\w* *: *$/g, 'except');
+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[];
+}

src/client/typeFormatters/elseFormatProvider.ts

@@ -0,0 +1,66 @@
+import { FormattingOptions, TextEdit, TextDocument } from 'vscode';
+import { Position, Range } from 'vscode';
+import { CodeBlockFormatProvider, BlockRegEx } from './contracts';
+import { IF_REGEX, ELIF_REGEX, ELSE_REGEX, FOR_IN_REGEX, ASYNC_FOR_IN_REGEX, WHILE_REGEX } from './contracts';
+import { DEF_REGEX, ASYNC_DEF_REGEX, CLASS_REGEX } from './contracts';
+
+export class ElseFormatProvider implements CodeBlockFormatProvider {
+    private regExps: BlockRegEx[];
+    private boundaryRegExps: BlockRegEx[];
+    constructor() {
+        this.regExps = [
+            IF_REGEX,
+            ELIF_REGEX,
+            FOR_IN_REGEX,
+            ASYNC_FOR_IN_REGEX,
+            WHILE_REGEX
+        ];
+        this.boundaryRegExps = [
+            DEF_REGEX,
+            ASYNC_DEF_REGEX,
+            CLASS_REGEX
+        ];
+    }
+    canProvideEdits(line: string): boolean {
+        return ELSE_REGEX.test(line);
+    }
+
+    provideEdits(document: TextDocument, position: Position, ch: string, options: FormattingOptions, line: string): TextEdit[] {
+        // We can have else for the following blocks:
+        // if:
+        // elif x:
+        // for x in y:
+        // while x:
+
+        const indexOfElse = line.indexOf(ELSE_REGEX.startWord);
+
+        // 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 line = document.lineAt(lineNumber).text;
+
+            // Oops, we've reached a boundary (like the function or class definition)
+            // Get out of here
+            if (this.boundaryRegExps.some(value => value.test(line))) {
+                return [];
+            }
+
+            const blockRegEx = this.regExps.find(value => value.test(line));
+            if (!blockRegEx) {
+                continue;
+            }
+
+            const startOfBlockInLine = line.indexOf(blockRegEx.startWord);
+            if (startOfBlockInLine > indexOfElse) {
+                continue;
+            }
+
+            const startPosition = new Position(position.line, 0);
+            const endPosition = new Position(position.line, indexOfElse - startOfBlockInLine);
+            return [
+                TextEdit.delete(new Range(startPosition, endPosition))
+            ];
+        }
+
+        return [];
+    }
+}