Merge pull request #473 from PowerShell/kapilmb/formatter-fixes

daviwil · web-flow · commit 3c7a723efac2 · 2017-01-30T13:54:53.000-08:00
Fixes some issues with code formatter
diff --git a/src/features/DocumentFormatter.ts b/src/features/DocumentFormatter.ts
@@ -2,18 +2,19 @@
  * Copyright (C) Microsoft Corporation. All rights reserved.
  *--------------------------------------------------------*/
 
+import * as path from "path";
 import vscode = require('vscode');
 import {
-    languages,
     TextDocument,
     TextEdit,
     FormattingOptions,
     CancellationToken,
     DocumentFormattingEditProvider,
     DocumentRangeFormattingEditProvider,
     Range,
+    TextEditor
 } from 'vscode';
-import { LanguageClient, RequestType, NotificationType } from 'vscode-languageclient';
+import { LanguageClient, RequestType } from 'vscode-languageclient';
 import Window = vscode.window;
 import { IFeature } from '../feature';
 import * as Settings from '../settings';
@@ -60,7 +61,7 @@ interface ScriptRegion {
 
 interface MarkerCorrection {
     name: string;
-    edits: ScriptRegion[]
+    edits: ScriptRegion[];
 }
 
 function editComparer(leftOperand: ScriptRegion, rightOperand: ScriptRegion): number {
@@ -81,7 +82,57 @@ function editComparer(leftOperand: ScriptRegion, rightOperand: ScriptRegion): nu
     }
 }
 
+class DocumentLocker {
+    private lockedDocuments: Object;
+
+    constructor() {
+        this.lockedDocuments = new Object();
+    }
+
+    isLocked(document: TextDocument): boolean {
+        return this.isLockedInternal(this.getKey(document));
+    }
+
+    lock(document: TextDocument, unlockWhenDone?: Thenable<any>): void {
+        this.lockInternal(this.getKey(document), unlockWhenDone);
+    }
+
+    unlock(document: TextDocument): void {
+        this.unlockInternal(this.getKey(document));
+    }
+
+    unlockAll(): void {
+        Object.keys(this.lockedDocuments).slice().forEach(documentKey => this.unlockInternal(documentKey));
+    }
+
+    private getKey(document: TextDocument): string {
+        return document.uri.toString();
+    }
+
+    private lockInternal(documentKey: string, unlockWhenDone?: Thenable<any>): void {
+        if (!this.isLockedInternal(documentKey)) {
+            this.lockedDocuments[documentKey] = true;
+        }
+
+        if (unlockWhenDone !== undefined) {
+            unlockWhenDone.then(() => this.unlockInternal(documentKey));
+        }
+    }
+
+    private unlockInternal(documentKey: string): void {
+        if (this.isLockedInternal(documentKey)) {
+            delete this.lockedDocuments[documentKey];
+        }
+    }
+
+    private isLockedInternal(documentKey: string): boolean {
+        return this.lockedDocuments.hasOwnProperty(documentKey);
+    }
+}
+
 class PSDocumentFormattingEditProvider implements DocumentFormattingEditProvider, DocumentRangeFormattingEditProvider {
+    private static documentLocker = new DocumentLocker();
+    private static statusBarTracker = new Object();
     private languageClient: LanguageClient;
 
     // The order in which the rules will be executed starting from the first element.
@@ -95,6 +146,10 @@ class PSDocumentFormattingEditProvider implements DocumentFormattingEditProvider
     // hence we keep this as an option but set it true by default.
     private aggregateUndoStop: boolean;
 
+    private get emptyPromise(): Promise<TextEdit[]> {
+        return Promise.resolve(TextEdit[0]);
+    }
+
     constructor(aggregateUndoStop = true) {
         this.aggregateUndoStop = aggregateUndoStop;
     }
@@ -112,19 +167,54 @@ class PSDocumentFormattingEditProvider implements DocumentFormattingEditProvider
         options: FormattingOptions,
         token: CancellationToken): TextEdit[] | Thenable<TextEdit[]> {
 
-        let textEdits: Thenable<TextEdit[]> = this.executeRulesInOrder(document, range, options, 0);
-        AnimatedStatusBar.showAnimatedStatusBarMessage("Formatting PowerShell document", textEdits);
+        let editor: TextEditor = this.getEditor(document);
+        if (editor === undefined) {
+            return this.emptyPromise;
+        }
+
+        // Check if the document is already being formatted.
+        // If so, then ignore the formatting request.
+        if (this.isDocumentLocked(document)) {
+            return this.emptyPromise;
+        }
+
+        let textEdits: Thenable<TextEdit[]> = this.executeRulesInOrder(editor, range, options, 0);
+        this.lockDocument(document, textEdits);
+        PSDocumentFormattingEditProvider.showStatusBar(document, textEdits);
         return textEdits;
     }
 
-    executeRulesInOrder(
-        document: TextDocument,
+    setLanguageClient(languageClient: LanguageClient): void {
+        this.languageClient = languageClient;
+
+        // setLanguageClient is called while restarting a session,
+        // so this makes sure we clean up the document locker and
+        // any residual status bars
+        PSDocumentFormattingEditProvider.documentLocker.unlockAll();
+        PSDocumentFormattingEditProvider.disposeAllStatusBars();
+    }
+
+    private getEditor(document: TextDocument): TextEditor {
+        return Window.visibleTextEditors.find((e, n, obj) => { return e.document === document; });
+    }
+
+    private isDocumentLocked(document: TextDocument): boolean {
+        return PSDocumentFormattingEditProvider.documentLocker.isLocked(document);
+    }
+
+    private lockDocument(document: TextDocument, unlockWhenDone: Thenable<any>): void {
+        PSDocumentFormattingEditProvider.documentLocker.lock(document, unlockWhenDone);
+    }
+
+    private executeRulesInOrder(
+        editor: TextEditor,
         range: Range,
         options: FormattingOptions,
         index: number): Thenable<TextEdit[]> {
         if (this.languageClient !== null && index < this.ruleOrder.length) {
-            let rule = this.ruleOrder[index];
+            let rule: string = this.ruleOrder[index];
             let uniqueEdits: ScriptRegion[] = [];
+            let document: TextDocument = editor.document;
             let edits: ScriptRegion[];
 
             return this.languageClient.sendRequest(
@@ -165,22 +255,29 @@ class PSDocumentFormattingEditProvider implements DocumentFormattingEditProvider
                     // we do not return a valid array because our text edits
                     // need to be executed in a particular order and it is
                     // easier if we perform the edits ourselves
-                    return this.applyEdit(uniqueEdits, range, 0, index);
+                    return this.applyEdit(editor, uniqueEdits, range, 0, index);
                 })
                 .then(() => {
                     // execute the same rule again if we left out violations
                     // on the same line
+                    let newIndex: number = index + 1;
                     if (uniqueEdits.length !== edits.length) {
-                        return this.executeRulesInOrder(document, range, options, index);
+                        newIndex = index;
                     }
-                    return this.executeRulesInOrder(document, range, options, index + 1);
+
+                    return this.executeRulesInOrder(editor, range, options, newIndex);
                 });
         } else {
-            return Promise.resolve(TextEdit[0]);
+            return this.emptyPromise;
         }
     }
 
-    applyEdit(edits: ScriptRegion[], range: Range, markerIndex: number, ruleIndex: number): Thenable<void> {
+    private applyEdit(
+        editor: TextEditor,
+        edits: ScriptRegion[],
+        range: Range,
+        markerIndex: number,
+        ruleIndex: number): Thenable<void> {
         if (markerIndex >= edits.length) {
             return;
         }
@@ -194,7 +291,7 @@ class PSDocumentFormattingEditProvider implements DocumentFormattingEditProvider
             edit.endLineNumber - 1,
             edit.endColumnNumber - 1);
         if (range === null || range.contains(editRange)) {
-            return Window.activeTextEditor.edit((editBuilder) => {
+            return editor.edit((editBuilder) => {
                 editBuilder.replace(
                     editRange,
                     edit.text);
@@ -203,15 +300,15 @@ class PSDocumentFormattingEditProvider implements DocumentFormattingEditProvider
                     undoStopAfter: undoStopAfter,
                     undoStopBefore: undoStopBefore
                 }).then((isEditApplied) => {
-                    return this.applyEdit(edits, range, markerIndex + 1, ruleIndex);
+                    return this.applyEdit(editor, edits, range, markerIndex + 1, ruleIndex);
                 }); // TODO handle rejection
         }
         else {
-            return this.applyEdit(edits, range, markerIndex + 1, ruleIndex);
+            return this.applyEdit(editor, edits, range, markerIndex + 1, ruleIndex);
         }
     }
 
-    getSelectionRange(document: TextDocument): Range {
+    private getSelectionRange(document: TextDocument): Range {
         let editor = vscode.window.visibleTextEditors.find(editor => editor.document === document);
         if (editor !== undefined) {
             return editor.selection as Range;
@@ -220,11 +317,7 @@ class PSDocumentFormattingEditProvider implements DocumentFormattingEditProvider
         return null;
     }
 
-    setLanguageClient(languageClient: LanguageClient): void {
-        this.languageClient = languageClient;
-    }
-
-    getSettings(rule: string): any {
+    private getSettings(rule: string): any {
         let psSettings: Settings.ISettings = Settings.load(Utils.PowerShellLanguageId);
         let ruleSettings = new Object();
         ruleSettings["Enable"] = true;
@@ -247,6 +340,25 @@ class PSDocumentFormattingEditProvider implements DocumentFormattingEditProvider
         settings[rule] = ruleSettings;
         return settings;
     }
+
+    private static showStatusBar(document: TextDocument, hideWhenDone: Thenable<any>): void {
+        let statusBar = AnimatedStatusBar.showAnimatedStatusBarMessage("Formatting PowerShell document", hideWhenDone);
+        this.statusBarTracker[document.uri.toString()] = statusBar;
+        hideWhenDone.then(() => {
+            this.disposeStatusBar(document.uri.toString());
+        });
+    }
+
+    private static disposeStatusBar(documentUri: string) {
+        if (this.statusBarTracker.hasOwnProperty(documentUri)) {
+            this.statusBarTracker[documentUri].dispose();
+            delete this.statusBarTracker[documentUri];
+        }
+    }
+
+    private static disposeAllStatusBars() {
+        Object.keys(this.statusBarTracker).slice().forEach((key) => this.disposeStatusBar(key));
+    }
 }
 
 export class DocumentFormatterFeature implements IFeature {
