Handle reloading of REPL Window (#24148)

anthonykim1 · web-flow · commit 8ac716d44833 · 2024-11-19T14:17:47.000-08:00
Resolves: #24021
diff --git a/src/client/repl/nativeRepl.ts b/src/client/repl/nativeRepl.ts
@@ -6,6 +6,7 @@ import {
     NotebookDocument,
     QuickPickItem,
     TextEditor,
+    Uri,
     workspace,
     WorkspaceFolder,
 } from 'vscode';
@@ -21,8 +22,11 @@ import { EventName } from '../telemetry/constants';
 import { sendTelemetryEvent } from '../telemetry';
 import { VariablesProvider } from './variables/variablesProvider';
 import { VariableRequester } from './variables/variableRequester';
+import { getTabNameForUri } from './replUtils';
+import { getWorkspaceStateValue, updateWorkspaceStateValue } from '../common/persistentState';
 
-let nativeRepl: NativeRepl | undefined; // In multi REPL scenario, hashmap of URI to Repl.
+export const NATIVE_REPL_URI_MEMENTO = 'nativeReplUri';
+let nativeRepl: NativeRepl | undefined;
 export class NativeRepl implements Disposable {
     // Adding ! since it will get initialized in create method, not the constructor.
     private pythonServer!: PythonServer;
@@ -65,10 +69,11 @@ export class NativeRepl implements Disposable {
      */
     private watchNotebookClosed(): void {
         this.disposables.push(
-            workspace.onDidCloseNotebookDocument((nb) => {
+            workspace.onDidCloseNotebookDocument(async (nb) => {
                 if (this.notebookDocument && nb.uri.toString() === this.notebookDocument.uri.toString()) {
                     this.notebookDocument = undefined;
                     this.newReplSession = true;
+                    await updateWorkspaceStateValue<string | undefined>(NATIVE_REPL_URI_MEMENTO, undefined);
                 }
             }),
         );
@@ -145,15 +150,37 @@ export class NativeRepl implements Disposable {
     /**
      * Function that opens interactive repl, selects kernel, and send/execute code to the native repl.
      */
-    public async sendToNativeRepl(code?: string): Promise<void> {
-        const notebookEditor = await openInteractiveREPL(this.replController, this.notebookDocument);
-        this.notebookDocument = notebookEditor.notebook;
-
-        if (this.notebookDocument) {
-            this.replController.updateNotebookAffinity(this.notebookDocument, NotebookControllerAffinity.Default);
-            await selectNotebookKernel(notebookEditor, this.replController.id, PVSC_EXTENSION_ID);
-            if (code) {
-                await executeNotebookCell(notebookEditor, code);
+    public async sendToNativeRepl(code?: string | undefined, preserveFocus: boolean = true): Promise<void> {
+        let wsMementoUri: Uri | undefined;
+
+        if (!this.notebookDocument) {
+            const wsMemento = getWorkspaceStateValue<string>(NATIVE_REPL_URI_MEMENTO);
+            wsMementoUri = wsMemento ? Uri.parse(wsMemento) : undefined;
+
+            if (!wsMementoUri || getTabNameForUri(wsMementoUri) !== 'Python REPL') {
+                await updateWorkspaceStateValue<string | undefined>(NATIVE_REPL_URI_MEMENTO, undefined);
+                wsMementoUri = undefined;
+            }
+        }
+
+        const notebookEditor = await openInteractiveREPL(
+            this.replController,
+            this.notebookDocument ?? wsMementoUri,
+            preserveFocus,
+        );
+        if (notebookEditor) {
+            this.notebookDocument = notebookEditor.notebook;
+            await updateWorkspaceStateValue<string | undefined>(
+                NATIVE_REPL_URI_MEMENTO,
+                this.notebookDocument.uri.toString(),
+            );
+
+            if (this.notebookDocument) {
+                this.replController.updateNotebookAffinity(this.notebookDocument, NotebookControllerAffinity.Default);
+                await selectNotebookKernel(notebookEditor, this.replController.id, PVSC_EXTENSION_ID);
+                if (code) {
+                    await executeNotebookCell(notebookEditor, code);
+                }
             }
         }
     }
diff --git a/src/client/repl/replCommandHandler.ts b/src/client/repl/replCommandHandler.ts
@@ -10,32 +10,48 @@ import {
     NotebookEdit,
     WorkspaceEdit,
     workspace,
+    Uri,
 } from 'vscode';
-import { getExistingReplViewColumn } from './replUtils';
+import { getExistingReplViewColumn, getTabNameForUri } from './replUtils';
 import { PVSC_EXTENSION_ID } from '../common/constants';
 
 /**
  * Function that opens/show REPL using IW UI.
  */
 export async function openInteractiveREPL(
     notebookController: NotebookController,
-    notebookDocument: NotebookDocument | undefined,
-): Promise<NotebookEditor> {
+    notebookDocument: NotebookDocument | Uri | undefined,
+    preserveFocus: boolean = true,
+): Promise<NotebookEditor | undefined> {
     let viewColumn = ViewColumn.Beside;
-
-    // Case where NotebookDocument (REPL document already exists in the tab)
-    if (notebookDocument) {
+    if (notebookDocument instanceof Uri) {
+        // Case where NotebookDocument is undefined, but workspace mementoURI exists.
+        notebookDocument = await workspace.openNotebookDocument(notebookDocument);
+    } else if (notebookDocument) {
+        // Case where NotebookDocument (REPL document already exists in the tab)
         const existingReplViewColumn = getExistingReplViewColumn(notebookDocument);
         viewColumn = existingReplViewColumn ?? viewColumn;
     } else if (!notebookDocument) {
-        // Case where NotebookDocument doesnt exist, create a blank one.
+        // Case where NotebookDocument doesnt exist, or
+        // became outdated (untitled.ipynb created without Python extension knowing, effectively taking over original Python REPL's URI)
         notebookDocument = await workspace.openNotebookDocument('jupyter-notebook');
     }
-    const editor = window.showNotebookDocument(notebookDocument!, {
+    const editor = await window.showNotebookDocument(notebookDocument!, {
         viewColumn,
         asRepl: 'Python REPL',
-        preserveFocus: true,
+        preserveFocus,
     });
+
+    // Sanity check that we opened a Native REPL from showNotebookDocument.
+    if (
+        !editor ||
+        !editor.notebook ||
+        !editor.notebook.uri ||
+        getTabNameForUri(editor.notebook.uri) !== 'Python REPL'
+    ) {
+        return undefined;
+    }
+
     await commands.executeCommand('notebook.selectKernel', {
         editor,
         id: notebookController.id,
diff --git a/src/client/repl/replCommands.ts b/src/client/repl/replCommands.ts
@@ -33,7 +33,7 @@ export async function registerStartNativeReplCommand(
             if (interpreter) {
                 if (interpreter) {
                     const nativeRepl = await getNativeRepl(interpreter, disposables);
-                    await nativeRepl.sendToNativeRepl();
+                    await nativeRepl.sendToNativeRepl(undefined, false);
                 }
             }
         }),
diff --git a/src/client/repl/replUtils.ts b/src/client/repl/replUtils.ts
@@ -99,3 +99,22 @@ export function getExistingReplViewColumn(notebookDocument: NotebookDocument): V
     }
     return undefined;
 }
+
+/**
+ * Function that will return tab name for before reloading VS Code
+ * This is so we can make sure tab name is still 'Python REPL' after reloading VS Code,
+ * and make sure Python REPL does not get 'merged' into unaware untitled.ipynb tab.
+ */
+export function getTabNameForUri(uri: Uri): string | undefined {
+    const tabGroups = window.tabGroups.all;
+
+    for (const tabGroup of tabGroups) {
+        for (const tab of tabGroup.tabs) {
+            if (tab.input instanceof TabInputNotebook && tab.input.uri.toString() === uri.toString()) {
+                return tab.label;
+            }
+        }
+    }
+
+    return undefined;
+}
diff --git a/src/test/repl/nativeRepl.test.ts b/src/test/repl/nativeRepl.test.ts
@@ -8,15 +8,17 @@ import { expect } from 'chai';
 import { IInterpreterService } from '../../client/interpreter/contracts';
 import { PythonEnvironment } from '../../client/pythonEnvironments/info';
 import { getNativeRepl, NativeRepl } from '../../client/repl/nativeRepl';
+import * as persistentState from '../../client/common/persistentState';
 
 suite('REPL - Native REPL', () => {
     let interpreterService: TypeMoq.IMock<IInterpreterService>;
 
     let disposable: TypeMoq.IMock<Disposable>;
     let disposableArray: Disposable[] = [];
-
     let setReplDirectoryStub: sinon.SinonStub;
     let setReplControllerSpy: sinon.SinonSpy;
+    let getWorkspaceStateValueStub: sinon.SinonStub;
+    let updateWorkspaceStateValueStub: sinon.SinonStub;
 
     setup(() => {
         interpreterService = TypeMoq.Mock.ofType<IInterpreterService>();
@@ -29,6 +31,7 @@ suite('REPL - Native REPL', () => {
         setReplDirectoryStub = sinon.stub(NativeRepl.prototype as any, 'setReplDirectory').resolves(); // Stubbing private method
         // Use a spy instead of a stub for setReplController
         setReplControllerSpy = sinon.spy(NativeRepl.prototype, 'setReplController');
+        updateWorkspaceStateValueStub = sinon.stub(persistentState, 'updateWorkspaceStateValue').resolves();
     });
 
     teardown(() => {
@@ -37,7 +40,6 @@ suite('REPL - Native REPL', () => {
                 d.dispose();
             }
         });
-
         disposableArray = [];
         sinon.restore();
     });
@@ -53,6 +55,32 @@ suite('REPL - Native REPL', () => {
         expect(createMethodStub.calledOnce).to.be.true;
     });
 
+    test('sendToNativeRepl should look for memento URI if notebook document is undefined', async () => {
+        getWorkspaceStateValueStub = sinon.stub(persistentState, 'getWorkspaceStateValue').returns(undefined);
+        interpreterService
+            .setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
+            .returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
+        const interpreter = await interpreterService.object.getActiveInterpreter();
+        const nativeRepl = await getNativeRepl(interpreter as PythonEnvironment, disposableArray);
+
+        nativeRepl.sendToNativeRepl(undefined, false);
+
+        expect(getWorkspaceStateValueStub.calledOnce).to.be.true;
+    });
+
+    test('sendToNativeRepl should call updateWorkspaceStateValue', async () => {
+        getWorkspaceStateValueStub = sinon.stub(persistentState, 'getWorkspaceStateValue').returns('myNameIsMemento');
+        interpreterService
+            .setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
+            .returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
+        const interpreter = await interpreterService.object.getActiveInterpreter();
+        const nativeRepl = await getNativeRepl(interpreter as PythonEnvironment, disposableArray);
+
+        nativeRepl.sendToNativeRepl(undefined, false);
+
+        expect(updateWorkspaceStateValueStub.calledOnce).to.be.true;
+    });
+
     test('create should call setReplDirectory, setReplController', async () => {
         const interpreter = await interpreterService.object.getActiveInterpreter();
         interpreterService
diff --git a/src/test/repl/replCommand.test.ts b/src/test/repl/replCommand.test.ts
@@ -24,6 +24,7 @@ suite('REPL - register native repl command', () => {
     let getNativeReplStub: sinon.SinonStub;
     let disposable: TypeMoq.IMock<Disposable>;
     let disposableArray: Disposable[] = [];
+
     setup(() => {
         interpreterService = TypeMoq.Mock.ofType<IInterpreterService>();
         commandManager = TypeMoq.Mock.ofType<ICommandManager>();

Original file line number	Diff line number	Diff line change
`@@ -33,7 +33,7 @@ export async function registerStartNativeReplCommand(`
`33`	`33`	`if (interpreter) {`
`34`	`34`	`if (interpreter) {`
`35`	`35`	`const nativeRepl = await getNativeRepl(interpreter, disposables);`
`36`		`- await nativeRepl.sendToNativeRepl();`
	`36`	`+ await nativeRepl.sendToNativeRepl(undefined, false);`
`37`	`37`	`}`
`38`	`38`	`}`
`39`	`39`	`}),`