Fix running Untitled files with the play button (#20955)

Author: Kartik Raj

src/client/terminals/codeExecution/codeExecutionManager.ts

@@ -91,11 +91,14 @@ export class CodeExecutionManager implements ICodeExecutionManager {
         sendTelemetryEvent(EventName.EXECUTION_CODE, undefined, { scope: 'file', trigger });
         const codeExecutionHelper = this.serviceContainer.get<ICodeExecutionHelper>(ICodeExecutionHelper);
         file = file instanceof Uri ? file : undefined;
-        const fileToExecute = file ? file : await codeExecutionHelper.getFileToExecute();
+        let fileToExecute = file ? file : await codeExecutionHelper.getFileToExecute();
         if (!fileToExecute) {
             return;
         }
-        await codeExecutionHelper.saveFileIfDirty(fileToExecute);
+        const fileAfterSave = await codeExecutionHelper.saveFileIfDirty(fileToExecute);
+        if (fileAfterSave) {
+            fileToExecute = fileAfterSave;
+        }
 
         try {
             const contents = await this.fileSystem.readFile(fileToExecute.fsPath);

src/client/terminals/codeExecution/helper.ts

@@ -5,7 +5,7 @@ import '../../common/extensions';
 import { inject, injectable } from 'inversify';
 import { l10n, Position, Range, TextEditor, Uri } from 'vscode';
 
-import { IApplicationShell, IDocumentManager } from '../../common/application/types';
+import { IApplicationShell, ICommandManager, IDocumentManager } from '../../common/application/types';
 import { PYTHON_LANGUAGE } from '../../common/constants';
 import * as internalScripts from '../../common/process/internal/scripts';
 import { IProcessServiceFactory } from '../../common/process/types';
@@ -14,6 +14,7 @@ import { IInterpreterService } from '../../interpreter/contracts';
 import { IServiceContainer } from '../../ioc/types';
 import { ICodeExecutionHelper } from '../types';
 import { traceError } from '../../logging';
+import { Resource } from '../../common/types';
 
 @injectable()
 export class CodeExecutionHelper implements ICodeExecutionHelper {
@@ -25,7 +26,7 @@ export class CodeExecutionHelper implements ICodeExecutionHelper {
 
     private readonly interpreterService: IInterpreterService;
 
-    constructor(@inject(IServiceContainer) serviceContainer: IServiceContainer) {
+    constructor(@inject(IServiceContainer) private readonly serviceContainer: IServiceContainer) {
         this.documentManager = serviceContainer.get<IDocumentManager>(IDocumentManager);
         this.applicationShell = serviceContainer.get<IApplicationShell>(IApplicationShell);
         this.processServiceFactory = serviceContainer.get<IProcessServiceFactory>(IProcessServiceFactory);
@@ -119,11 +120,17 @@ export class CodeExecutionHelper implements ICodeExecutionHelper {
         return code;
     }
 
-    public async saveFileIfDirty(file: Uri): Promise<void> {
+    public async saveFileIfDirty(file: Uri): Promise<Resource> {
         const docs = this.documentManager.textDocuments.filter((d) => d.uri.path === file.path);
         if (docs.length === 1 && docs[0].isDirty) {
-            await docs[0].save();
+            const deferred = createDeferred<Uri>();
+            this.documentManager.onDidSaveTextDocument((e) => deferred.resolve(e.uri));
+            const commandManager = this.serviceContainer.get<ICommandManager>(ICommandManager);
+            await commandManager.executeCommand('workbench.action.files.save', file);
+            const savedFileUri = await deferred.promise;
+            return savedFileUri;
         }
+        return undefined;
     }
 }
 

src/client/terminals/types.ts

@@ -2,7 +2,7 @@
 // Licensed under the MIT License.
 
 import { Event, Terminal, TextEditor, Uri } from 'vscode';
-import { IDisposable } from '../common/types';
+import { IDisposable, Resource } from '../common/types';
 
 export const ICodeExecutionService = Symbol('ICodeExecutionService');
 
@@ -17,7 +17,7 @@ export const ICodeExecutionHelper = Symbol('ICodeExecutionHelper');
 export interface ICodeExecutionHelper {
     normalizeLines(code: string): Promise<string>;
     getFileToExecute(): Promise<Uri | undefined>;
-    saveFileIfDirty(file: Uri): Promise<void>;
+    saveFileIfDirty(file: Uri): Promise<Resource>;
     getSelectedTextToExecute(textEditor: TextEditor): Promise<string | undefined>;
 }
 

src/test/terminals/codeExecution/helper.test.ts

@@ -8,8 +8,8 @@ import * as fs from 'fs-extra';
 import * as path from 'path';
 import { SemVer } from 'semver';
 import * as TypeMoq from 'typemoq';
-import { Position, Range, Selection, TextDocument, TextEditor, TextLine, Uri } from 'vscode';
-import { IApplicationShell, IDocumentManager } from '../../../client/common/application/types';
+import { EventEmitter, Position, Range, Selection, TextDocument, TextEditor, TextLine, Uri } from 'vscode';
+import { IApplicationShell, ICommandManager, IDocumentManager } from '../../../client/common/application/types';
 import { EXTENSION_ROOT_DIR, PYTHON_LANGUAGE } from '../../../client/common/constants';
 import '../../../client/common/extensions';
 import { ProcessService } from '../../../client/common/process/proc';
@@ -37,6 +37,7 @@ suite('Terminal - Code Execution Helper', () => {
     let editor: TypeMoq.IMock<TextEditor>;
     let processService: TypeMoq.IMock<IProcessService>;
     let interpreterService: TypeMoq.IMock<IInterpreterService>;
+    let commandManager: TypeMoq.IMock<ICommandManager>;
     const workingPython: PythonEnvironment = {
         path: PYTHON_PATH,
         version: new SemVer('3.6.6-final'),
@@ -49,6 +50,7 @@ suite('Terminal - Code Execution Helper', () => {
 
     setup(() => {
         const serviceContainer = TypeMoq.Mock.ofType<IServiceContainer>();
+        commandManager = TypeMoq.Mock.ofType<ICommandManager>();
         documentManager = TypeMoq.Mock.ofType<IDocumentManager>();
         applicationShell = TypeMoq.Mock.ofType<IApplicationShell>();
         const envVariablesProvider = TypeMoq.Mock.ofType<IEnvironmentVariablesProvider>();
@@ -79,6 +81,7 @@ suite('Terminal - Code Execution Helper', () => {
         serviceContainer
             .setup((c) => c.get(TypeMoq.It.isValue(IApplicationShell), TypeMoq.It.isAny()))
             .returns(() => applicationShell.object);
+        serviceContainer.setup((c) => c.get(TypeMoq.It.isValue(ICommandManager))).returns(() => commandManager.object);
         serviceContainer
             .setup((c) => c.get(TypeMoq.It.isValue(IEnvironmentVariablesProvider), TypeMoq.It.isAny()))
             .returns(() => envVariablesProvider.object);
@@ -364,15 +367,24 @@ suite('Terminal - Code Execution Helper', () => {
             .setup((d) => d.textDocuments)
             .returns(() => [document.object])
             .verifiable(TypeMoq.Times.once());
-        document.setup((doc) => doc.isUntitled).returns(() => false);
+        const saveEmitter = new EventEmitter<TextDocument>();
+        documentManager.setup((d) => d.onDidSaveTextDocument).returns(() => saveEmitter.event);
+        document.setup((doc) => doc.isUntitled).returns(() => true);
         document.setup((doc) => doc.isDirty).returns(() => true);
         document.setup((doc) => doc.languageId).returns(() => PYTHON_LANGUAGE);
-        const expectedUri = Uri.file('one.py');
-        document.setup((doc) => doc.uri).returns(() => expectedUri);
-
-        await helper.saveFileIfDirty(expectedUri);
-        documentManager.verifyAll();
-        document.verify((doc) => doc.save(), TypeMoq.Times.once());
+        const untitledUri = Uri.file('Untitled-1');
+        document.setup((doc) => doc.uri).returns(() => untitledUri);
+        const savedDocument = TypeMoq.Mock.ofType<TextDocument>();
+        const expectedSavedUri = Uri.file('one.py');
+        savedDocument.setup((doc) => doc.uri).returns(() => expectedSavedUri);
+        commandManager
+            .setup((c) => c.executeCommand('workbench.action.files.save', untitledUri))
+            .callback(() => saveEmitter.fire(savedDocument.object))
+            .returns(() => Promise.resolve());
+
+        const savedUri = await helper.saveFileIfDirty(untitledUri);
+
+        expect(savedUri?.fsPath).to.be.equal(expectedSavedUri.fsPath);
     });
 
     test('File will be not saved if file is not dirty', async () => {