From 8ac716d44833fd4e29739869827e0df6c2b0eba5 Mon Sep 17 00:00:00 2001 From: Anthony Kim <62267334+anthonykim1@users.noreply.github.com> Date: Tue, 19 Nov 2024 14:17:47 -0800 Subject: [PATCH] Handle reloading of REPL Window (#24148) Resolves: https://github.com/microsoft/vscode-python/issues/24021 --- src/client/repl/nativeRepl.ts | 49 +++++++++++++++++++++------ src/client/repl/replCommandHandler.ts | 34 ++++++++++++++----- src/client/repl/replCommands.ts | 2 +- src/client/repl/replUtils.ts | 19 +++++++++++ src/test/repl/nativeRepl.test.ts | 32 +++++++++++++++-- src/test/repl/replCommand.test.ts | 1 + 6 files changed, 114 insertions(+), 23 deletions(-) diff --git a/src/client/repl/nativeRepl.ts b/src/client/repl/nativeRepl.ts index 8304a27db6ea..6edd3cbd70a7 100644 --- 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(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 { - 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 { + let wsMementoUri: Uri | undefined; + + if (!this.notebookDocument) { + const wsMemento = getWorkspaceStateValue(NATIVE_REPL_URI_MEMENTO); + wsMementoUri = wsMemento ? Uri.parse(wsMemento) : undefined; + + if (!wsMementoUri || getTabNameForUri(wsMementoUri) !== 'Python REPL') { + await updateWorkspaceStateValue(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( + 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 index 13e7607b5cc8..89ccbe11c337 100644 --- a/src/client/repl/replCommandHandler.ts +++ b/src/client/repl/replCommandHandler.ts @@ -10,8 +10,9 @@ import { NotebookEdit, WorkspaceEdit, workspace, + Uri, } from 'vscode'; -import { getExistingReplViewColumn } from './replUtils'; +import { getExistingReplViewColumn, getTabNameForUri } from './replUtils'; import { PVSC_EXTENSION_ID } from '../common/constants'; /** @@ -19,23 +20,38 @@ import { PVSC_EXTENSION_ID } from '../common/constants'; */ export async function openInteractiveREPL( notebookController: NotebookController, - notebookDocument: NotebookDocument | undefined, -): Promise { + notebookDocument: NotebookDocument | Uri | undefined, + preserveFocus: boolean = true, +): Promise { 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 index eb2eb49aea76..f260185988a8 100644 --- 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 index 5f58f461155b..0c2c4ba0d84e 100644 --- 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 index 0fc55abe1a64..999bc656c64d 100644 --- 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; let disposable: TypeMoq.IMock; let disposableArray: Disposable[] = []; - let setReplDirectoryStub: sinon.SinonStub; let setReplControllerSpy: sinon.SinonSpy; + let getWorkspaceStateValueStub: sinon.SinonStub; + let updateWorkspaceStateValueStub: sinon.SinonStub; setup(() => { interpreterService = TypeMoq.Mock.ofType(); @@ -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 index 444b8e5f16b9..7c26ebd69c80 100644 --- 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; let disposableArray: Disposable[] = []; + setup(() => { interpreterService = TypeMoq.Mock.ofType(); commandManager = TypeMoq.Mock.ofType();