Fix colon-triggered block formatting

news/2 Fixes/2714.md

@@ -0,0 +1 @@
+Fix colon-triggered block formatting

src/client/extension.ts

@@ -59,6 +59,7 @@ import { EDITOR_LOAD } from './telemetry/constants';
 import { registerTypes as commonRegisterTerminalTypes } from './terminals/serviceRegistry';
 import { ICodeExecutionManager, ITerminalAutoActivation } from './terminals/types';
 import { BlockFormatProviders } from './typeFormatters/blockFormatProvider';
+import { OnTypeFormattingDispatcher } from './typeFormatters/dispatcher';
 import { OnEnterFormatter } from './typeFormatters/onEnterFormatter';
 import { TEST_OUTPUT_CHANNEL } from './unittests/common/constants';
 import { registerTypes as unitTestsRegisterTypes } from './unittests/serviceRegistry';
@@ -136,8 +137,14 @@ export async function activate(context: ExtensionContext): Promise<IExtensionApi
         context.subscriptions.push(languages.registerDocumentRangeFormattingEditProvider(PYTHON, formatProvider));
     }
 
-    context.subscriptions.push(languages.registerOnTypeFormattingEditProvider(PYTHON, new BlockFormatProviders(), ':'));
-    context.subscriptions.push(languages.registerOnTypeFormattingEditProvider(PYTHON, new OnEnterFormatter(), '\n'));
+    const onTypeDispatcher = new OnTypeFormattingDispatcher({
+        '\n': new OnEnterFormatter(),
+        ':': new BlockFormatProviders()
+    });
+    const onTypeTriggers = onTypeDispatcher.getTriggerCharacters();
+    if (onTypeTriggers) {
+        context.subscriptions.push(languages.registerOnTypeFormattingEditProvider(PYTHON, onTypeDispatcher, onTypeTriggers.first, ...onTypeTriggers.more));
+    }
 
     const deprecationMgr = serviceContainer.get<IFeatureDeprecationManager>(IFeatureDeprecationManager);
     deprecationMgr.initialize();

src/client/typeFormatters/dispatcher.ts

@@ -0,0 +1,40 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+import { CancellationToken, FormattingOptions, OnTypeFormattingEditProvider, Position, ProviderResult, TextDocument, TextEdit } from 'vscode';
+
+export class OnTypeFormattingDispatcher implements OnTypeFormattingEditProvider {
+    private readonly providers: { [key: string]: OnTypeFormattingEditProvider };
+
+    constructor(providers: { [key: string]: OnTypeFormattingEditProvider }) {
+        this.providers = providers;
+    }
+
+    public provideOnTypeFormattingEdits(document: TextDocument, position: Position, ch: string, options: FormattingOptions, cancellationToken: CancellationToken): ProviderResult<TextEdit[]> {
+        const provider = this.providers[ch];
+
+        if (provider) {
+            return provider.provideOnTypeFormattingEdits(document, position, ch, options, cancellationToken);
+        }
+
+        return [];
+    }
+
+    public getTriggerCharacters(): { first: string; more: string[] } | undefined {
+        let keys = Object.keys(this.providers);
+        keys = keys.sort(); // Make output deterministic
+
+        const first = keys.shift();
+
+        if (first) {
+            return {
+                first: first,
+                more: keys
+            };
+        }
+
+        return undefined;
+    }
+}

src/test/format/extension.dispatch.test.ts

@@ -0,0 +1,81 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+import * as assert from 'assert';
+import * as TypeMoq from 'typemoq';
+import { CancellationToken, FormattingOptions, OnTypeFormattingEditProvider, Position, ProviderResult, TextDocument, TextEdit } from 'vscode';
+import { OnTypeFormattingDispatcher } from '../../client/typeFormatters/dispatcher';
+
+suite('Formatting - Dispatcher', () => {
+    const doc = TypeMoq.Mock.ofType<TextDocument>();
+    const pos = TypeMoq.Mock.ofType<Position>();
+    const opt = TypeMoq.Mock.ofType<FormattingOptions>();
+    const token = TypeMoq.Mock.ofType<CancellationToken>();
+    const edits = TypeMoq.Mock.ofType<ProviderResult<TextEdit[]>>();
+
+    test('No providers', async () => {
+        const dispatcher = new OnTypeFormattingDispatcher({});
+
+        const triggers = dispatcher.getTriggerCharacters();
+        assert.equal(triggers, undefined, 'Trigger was not undefined');
+
+        const result = await dispatcher.provideOnTypeFormattingEdits(doc.object, pos.object, '\n', opt.object, token.object);
+        assert.deepStrictEqual(result, [], 'Did not return an empty list of edits');
+    });
+
+    test('Single provider', () => {
+        const provider = setupProvider(doc.object, pos.object, ':', opt.object, token.object, edits.object);
+
+        const dispatcher = new OnTypeFormattingDispatcher({
+            ':': provider.object
+        });
+
+        const triggers = dispatcher.getTriggerCharacters();
+        assert.deepStrictEqual(triggers, { first: ':', more: [] }, 'Did not return correct triggers');
+
+        const result = dispatcher.provideOnTypeFormattingEdits(doc.object, pos.object, ':', opt.object, token.object);
+        assert.equal(result, edits.object, 'Did not return correct edits');
+
+        provider.verifyAll();
+    });
+
+    test('Two providers', () => {
+        const colonProvider = setupProvider(doc.object, pos.object, ':', opt.object, token.object, edits.object);
+
+        const doc2 = TypeMoq.Mock.ofType<TextDocument>();
+        const pos2 = TypeMoq.Mock.ofType<Position>();
+        const opt2 = TypeMoq.Mock.ofType<FormattingOptions>();
+        const token2 = TypeMoq.Mock.ofType<CancellationToken>();
+        const edits2 = TypeMoq.Mock.ofType<ProviderResult<TextEdit[]>>();
+
+        const newlineProvider = setupProvider(doc2.object, pos2.object, '\n', opt2.object, token2.object, edits2.object);
+
+        const dispatcher = new OnTypeFormattingDispatcher({
+            ':': colonProvider.object,
+            '\n': newlineProvider.object
+        });
+
+        const triggers = dispatcher.getTriggerCharacters();
+        assert.deepStrictEqual(triggers, { first: '\n', more: [':'] }, 'Did not return correct triggers');
+
+        const result = dispatcher.provideOnTypeFormattingEdits(doc.object, pos.object, ':', opt.object, token.object);
+        assert.equal(result, edits.object, 'Did not return correct editsfor colon provider');
+
+        const result2 = dispatcher.provideOnTypeFormattingEdits(doc2.object, pos2.object, '\n', opt2.object, token2.object);
+        assert.equal(result2, edits2.object, 'Did not return correct edits for newline provider');
+
+        colonProvider.verifyAll();
+        newlineProvider.verifyAll();
+    });
+
+    function setupProvider(document: TextDocument, position: Position, ch: string, options: FormattingOptions, cancellationToken: CancellationToken,
+        result: ProviderResult<TextEdit[]>): TypeMoq.IMock<OnTypeFormattingEditProvider> {
+        const provider = TypeMoq.Mock.ofType<OnTypeFormattingEditProvider>();
+        provider.setup(p => p.provideOnTypeFormattingEdits(document, position, ch, options, cancellationToken))
+            .returns(() => result)
+            .verifiable(TypeMoq.Times.once());
+        return provider;
+    }
+});