Fix native repl not using env extension

python_files/vscode_pytest/__init__.py

@@ -190,7 +190,7 @@ def pytest_exception_interact(node, call, report):
             send_execution_message(
                 os.fsdecode(cwd),
                 "success",
-                collected_test if collected_test else None,
+                collected_test or None,
             )
 
 
@@ -314,7 +314,7 @@ def pytest_report_teststatus(report, config):  # noqa: ARG001
             send_execution_message(
                 os.fsdecode(cwd),
                 "success",
-                collected_test if collected_test else None,
+                collected_test or None,
             )
     yield
 
@@ -348,7 +348,7 @@ def pytest_runtest_protocol(item, nextitem):  # noqa: ARG001
             send_execution_message(
                 os.fsdecode(cwd),
                 "success",
-                collected_test if collected_test else None,
+                collected_test or None,
             )
     yield
 
@@ -1024,7 +1024,7 @@ def get_node_path(
         except Exception as e:
             raise VSCodePytestError(
                 f"Error occurred while calculating symlink equivalent from node path: {e}"
-                f"\n SYMLINK_PATH: {SYMLINK_PATH}, \n node path: {node_path}, \n cwd: {_CACHED_CWD if _CACHED_CWD else pathlib.Path.cwd()}"
+                f"\n SYMLINK_PATH: {SYMLINK_PATH}, \n node path: {node_path}, \n cwd: {_CACHED_CWD or pathlib.Path.cwd()}"
             ) from e
     else:
         result = node_path

src/client/repl/nativeRepl.ts

@@ -4,9 +4,10 @@
 // Native Repl class that holds instance of pythonServer and replController
 
 import { NotebookController, NotebookDocument, QuickPickItem, TextEditor, Uri, WorkspaceFolder } from 'vscode';
+import * as path from 'path';
 import { Disposable } from 'vscode-jsonrpc';
 import { PVSC_EXTENSION_ID } from '../common/constants';
-import { showQuickPick } from '../common/vscodeApis/windowApis';
+import { showNotebookDocument, showQuickPick } from '../common/vscodeApis/windowApis';
 import { getWorkspaceFolders, onDidCloseNotebookDocument } from '../common/vscodeApis/workspaceApis';
 import { PythonEnvironment } from '../pythonEnvironments/info';
 import { createPythonServer, PythonServer } from './pythonServer';
@@ -18,6 +19,8 @@ import { VariablesProvider } from './variables/variablesProvider';
 import { VariableRequester } from './variables/variableRequester';
 import { getTabNameForUri } from './replUtils';
 import { getWorkspaceStateValue, updateWorkspaceStateValue } from '../common/persistentState';
+import { onDidChangeEnvironmentEnvExt, useEnvExtension } from '../envExt/api.internal';
+import { getActiveInterpreterLegacy } from '../envExt/api.legacy';
 
 export const NATIVE_REPL_URI_MEMENTO = 'nativeReplUri';
 let nativeRepl: NativeRepl | undefined;
@@ -37,6 +40,10 @@ export class NativeRepl implements Disposable {
 
     public newReplSession: boolean | undefined = true;
 
+    private envChangeListenerRegistered = false;
+
+    private pendingInterpreterChangeTimer?: NodeJS.Timeout;
+
     // TODO: In the future, could also have attribute of URI for file specific REPL.
     private constructor() {
         this.watchNotebookClosed();
@@ -48,12 +55,17 @@ export class NativeRepl implements Disposable {
         nativeRepl.interpreter = interpreter;
         await nativeRepl.setReplDirectory();
         nativeRepl.pythonServer = createPythonServer([interpreter.path as string], nativeRepl.cwd);
+        nativeRepl.disposables.push(nativeRepl.pythonServer);
         nativeRepl.setReplController();
+        nativeRepl.registerInterpreterChangeHandler();
 
         return nativeRepl;
     }
 
     dispose(): void {
+        if (this.pendingInterpreterChangeTimer) {
+            clearTimeout(this.pendingInterpreterChangeTimer);
+        }
         this.disposables.forEach((d) => d.dispose());
     }
 
@@ -116,8 +128,8 @@ export class NativeRepl implements Disposable {
     /**
      * Function that check if NotebookController for REPL exists, and returns it in Singleton manner.
      */
-    public setReplController(): NotebookController {
-        if (!this.replController) {
+    public setReplController(force: boolean = false): NotebookController {
+        if (!this.replController || force) {
             this.replController = createReplController(this.interpreter!.path, this.disposables, this.cwd);
             this.replController.variableProvider = new VariablesProvider(
                 new VariableRequester(this.pythonServer),
@@ -128,6 +140,65 @@ export class NativeRepl implements Disposable {
         return this.replController;
     }
 
+    private registerInterpreterChangeHandler(): void {
+        if (!useEnvExtension() || this.envChangeListenerRegistered) {
+            return;
+        }
+        this.envChangeListenerRegistered = true;
+        this.disposables.push(
+            onDidChangeEnvironmentEnvExt((event) => {
+                this.updateInterpreterForChange(event.uri).catch(() => undefined);
+            }),
+        );
+    }
+
+    private async updateInterpreterForChange(resource?: Uri): Promise<void> {
+        if (this.pythonServer?.isExecuting) {
+            this.scheduleInterpreterUpdate(resource);
+            return;
+        }
+        if (!this.shouldApplyInterpreterChange(resource)) {
+            return;
+        }
+        const scope = resource ?? (this.cwd ? Uri.file(this.cwd) : undefined);
+        const interpreter = await getActiveInterpreterLegacy(scope);
+        if (!interpreter || interpreter.path === this.interpreter?.path) {
+            return;
+        }
+
+        this.interpreter = interpreter;
+        this.pythonServer.dispose();
+        this.pythonServer = createPythonServer([interpreter.path as string], this.cwd);
+        this.disposables.push(this.pythonServer);
+        if (this.replController) {
+            this.replController.dispose();
+        }
+        this.setReplController(true);
+
+        if (this.notebookDocument) {
+            const notebookEditor = await showNotebookDocument(this.notebookDocument, { preserveFocus: true });
+            await selectNotebookKernel(notebookEditor, this.replController.id, PVSC_EXTENSION_ID);
+        }
+    }
+
+    private scheduleInterpreterUpdate(resource?: Uri): void {
+        if (this.pendingInterpreterChangeTimer) {
+            clearTimeout(this.pendingInterpreterChangeTimer);
+        }
+        this.pendingInterpreterChangeTimer = setTimeout(() => {
+            this.pendingInterpreterChangeTimer = undefined;
+            this.updateInterpreterForChange(resource).catch(() => undefined);
+        }, 200);
+    }
+
+    private shouldApplyInterpreterChange(resource?: Uri): boolean {
+        if (!resource || !this.cwd) {
+            return true;
+        }
+        const relative = path.relative(this.cwd, resource.fsPath);
+        return relative === '' || (!relative.startsWith('..') && !path.isAbsolute(relative));
+    }
+
     /**
      * Function that checks if native REPL's text input box contains complete code.
      * @returns Promise<boolean> - True if complete/Valid code is present, False otherwise.

src/client/repl/pythonServer.ts

@@ -16,6 +16,8 @@ export interface ExecutionResult {
 
 export interface PythonServer extends Disposable {
     onCodeExecuted: Event<void>;
+    readonly isExecuting: boolean;
+    readonly isDisposed: boolean;
     execute(code: string): Promise<ExecutionResult | undefined>;
     executeSilently(code: string): Promise<ExecutionResult | undefined>;
     interrupt(): void;
@@ -30,6 +32,18 @@ class PythonServerImpl implements PythonServer, Disposable {
 
     onCodeExecuted = this._onCodeExecuted.event;
 
+    private inFlightRequests = 0;
+
+    private disposed = false;
+
+    public get isExecuting(): boolean {
+        return this.inFlightRequests > 0;
+    }
+
+    public get isDisposed(): boolean {
+        return this.disposed;
+    }
+
     constructor(private connection: rpc.MessageConnection, private pythonServer: ch.ChildProcess) {
         this.initialize();
         this.input();
@@ -41,6 +55,14 @@ class PythonServerImpl implements PythonServer, Disposable {
                 traceLog('Log:', message);
             }),
         );
+        this.pythonServer.on('exit', (code) => {
+            traceError(`Python server exited with code ${code}`);
+            this.markDisposed();
+        });
+        this.pythonServer.on('error', (err) => {
+            traceError(err);
+            this.markDisposed();
+        });
         this.connection.listen();
     }
 
@@ -75,12 +97,15 @@ class PythonServerImpl implements PythonServer, Disposable {
     }
 
     private async executeCode(code: string): Promise<ExecutionResult | undefined> {
+        this.inFlightRequests += 1;
         try {
             const result = await this.connection.sendRequest('execute', code);
             return result as ExecutionResult;
         } catch (err) {
             const error = err as Error;
             traceError(`Error getting response from REPL server:`, error);
+        } finally {
+            this.inFlightRequests -= 1;
         }
         return undefined;
     }
@@ -93,39 +118,47 @@ class PythonServerImpl implements PythonServer, Disposable {
     }
 
     public async checkValidCommand(code: string): Promise<boolean> {
-        const completeCode: ExecutionResult = await this.connection.sendRequest('check_valid_command', code);
-        if (completeCode.output === 'True') {
-            return new Promise((resolve) => resolve(true));
+        this.inFlightRequests += 1;
+        try {
+            const completeCode: ExecutionResult = await this.connection.sendRequest('check_valid_command', code);
+            return completeCode.output === 'True';
+        } finally {
+            this.inFlightRequests -= 1;
         }
-        return new Promise((resolve) => resolve(false));
     }
 
     public dispose(): void {
+        if (this.disposed) {
+            return;
+        }
+        this.disposed = true;
         this.connection.sendNotification('exit');
         this.disposables.forEach((d) => d.dispose());
         this.connection.dispose();
         serverInstance = undefined;
     }
+
+    private markDisposed(): void {
+        if (this.disposed) {
+            return;
+        }
+        this.disposed = true;
+        this.connection.dispose();
+        serverInstance = undefined;
+    }
 }
 
 export function createPythonServer(interpreter: string[], cwd?: string): PythonServer {
-    if (serverInstance) {
+    if (serverInstance && !serverInstance.isDisposed) {
         return serverInstance;
     }
 
     const pythonServer = ch.spawn(interpreter[0], [...interpreter.slice(1), SERVER_PATH], {
         cwd, // Launch with correct workspace directory
     });
-
     pythonServer.stderr.on('data', (data) => {
         traceError(data.toString());
     });
-    pythonServer.on('exit', (code) => {
-        traceError(`Python server exited with code ${code}`);
-    });
-    pythonServer.on('error', (err) => {
-        traceError(err);
-    });
     const connection = rpc.createMessageConnection(
         new rpc.StreamMessageReader(pythonServer.stdout),
         new rpc.StreamMessageWriter(pythonServer.stdin),

src/client/repl/replCommands.ts

@@ -2,7 +2,6 @@ import { commands, Uri, window } from 'vscode';
 import { Disposable } from 'vscode-jsonrpc';
 import { ICommandManager } from '../common/application/types';
 import { Commands } from '../common/constants';
-import { noop } from '../common/utils/misc';
 import { IInterpreterService } from '../interpreter/contracts';
 import { ICodeExecutionHelper } from '../terminals/types';
 import { getNativeRepl } from './nativeRepl';
@@ -102,14 +101,13 @@ export async function registerReplExecuteOnEnter(
 }
 
 async function onInputEnter(
-    uri: Uri,
+    uri: Uri | undefined,
     commandName: string,
     interpreterService: IInterpreterService,
     disposables: Disposable[],
 ): Promise<void> {
-    const interpreter = await interpreterService.getActiveInterpreter(uri);
+    const interpreter = await getActiveInterpreter(uri, interpreterService);
     if (!interpreter) {
-        commands.executeCommand(Commands.TriggerEnvironmentSelection, uri).then(noop, noop);
         return;
     }
 

src/client/repl/replUtils.ts

@@ -66,12 +66,13 @@ export function isMultiLineText(textEditor: TextEditor): boolean {
  * Function will also return undefined or active interpreter
  */
 export async function getActiveInterpreter(
-    uri: Uri,
+    uri: Uri | undefined,
     interpreterService: IInterpreterService,
 ): Promise<PythonEnvironment | undefined> {
-    const interpreter = await interpreterService.getActiveInterpreter(uri);
+    const resource = uri ?? getActiveResource();
+    const interpreter = await interpreterService.getActiveInterpreter(resource);
     if (!interpreter) {
-        commands.executeCommand(Commands.TriggerEnvironmentSelection, uri).then(noop, noop);
+        commands.executeCommand(Commands.TriggerEnvironmentSelection, resource).then(noop, noop);
         return undefined;
     }
     return interpreter;

src/test/repl/nativeRepl.test.ts

@@ -129,6 +129,8 @@ suite('REPL - Native REPL', () => {
             input: sinon.stub(),
             checkValidCommand: sinon.stub().resolves(true),
             dispose: sinon.stub(),
+            isExecuting: false,
+            isDisposed: false,
         };
 
         // Track the number of times createPythonServer was called

src/test/repl/replCommand.test.ts

@@ -1,6 +1,6 @@
 // Create test suite and test cases for the `replUtils` module
 import * as TypeMoq from 'typemoq';
-import { Disposable } from 'vscode';
+import { commands, Disposable, Uri } from 'vscode';
 import * as sinon from 'sinon';
 import { expect } from 'chai';
 import { IInterpreterService } from '../../client/interpreter/contracts';
@@ -9,6 +9,7 @@ import { ICodeExecutionHelper } from '../../client/terminals/types';
 import * as replCommands from '../../client/repl/replCommands';
 import * as replUtils from '../../client/repl/replUtils';
 import * as nativeRepl from '../../client/repl/nativeRepl';
+import * as windowApis from '../../client/common/vscodeApis/windowApis';
 import { Commands } from '../../client/common/constants';
 import { PythonEnvironment } from '../../client/pythonEnvironments/info';
 
@@ -203,3 +204,47 @@ suite('REPL - register native repl command', () => {
         sinon.assert.notCalled(getNativeReplStub);
     });
 });
+
+suite('Native REPL getActiveInterpreter', () => {
+    let interpreterService: TypeMoq.IMock<IInterpreterService>;
+    let executeCommandStub: sinon.SinonStub;
+    let getActiveResourceStub: sinon.SinonStub;
+
+    setup(() => {
+        interpreterService = TypeMoq.Mock.ofType<IInterpreterService>();
+        executeCommandStub = sinon.stub(commands, 'executeCommand').resolves(undefined);
+        getActiveResourceStub = sinon.stub(windowApis, 'getActiveResource');
+    });
+
+    teardown(() => {
+        sinon.restore();
+    });
+
+    test('Uses active resource when uri is undefined', async () => {
+        const resource = Uri.file('/workspace/app.py');
+        const expected = ({ path: 'ps' } as unknown) as PythonEnvironment;
+        getActiveResourceStub.returns(resource);
+        interpreterService
+            .setup((i) => i.getActiveInterpreter(TypeMoq.It.isValue(resource)))
+            .returns(() => Promise.resolve(expected));
+
+        const result = await replUtils.getActiveInterpreter(undefined, interpreterService.object);
+
+        expect(result).to.equal(expected);
+        interpreterService.verify((i) => i.getActiveInterpreter(TypeMoq.It.isValue(resource)), TypeMoq.Times.once());
+        sinon.assert.notCalled(executeCommandStub);
+    });
+
+    test('Triggers environment selection using active resource when interpreter is missing', async () => {
+        const resource = Uri.file('/workspace/app.py');
+        getActiveResourceStub.returns(resource);
+        interpreterService
+            .setup((i) => i.getActiveInterpreter(TypeMoq.It.isValue(resource)))
+            .returns(() => Promise.resolve(undefined));
+
+        const result = await replUtils.getActiveInterpreter(undefined, interpreterService.object);
+
+        expect(result).to.equal(undefined);
+        sinon.assert.calledWith(executeCommandStub, Commands.TriggerEnvironmentSelection, resource);
+    });
+});