Adding PYTHONSTARTUP with shell integration to environment variable collection

package.json

@@ -631,6 +631,12 @@
                     "scope": "resource",
                     "type": "array"
                 },
+                "python.terminal.enable Shell Integration in Python Terminal REPL": {
+                    "default": true,
+                    "description": "%python.terminal.enableSITerminal.description%",
+                    "scope": "resource",
+                    "type": "boolean"
+                },
                 "python.REPL.enableREPLSmartSend": {
                     "default": true,
                     "description": "%python.EnableREPLSmartSend.description%",

package.nls.json

@@ -68,6 +68,7 @@
     "python.tensorBoard.logDirectory.description": "Set this setting to your preferred TensorBoard log directory to skip log directory prompt when starting TensorBoard.",
     "python.tensorBoard.logDirectory.markdownDeprecationMessage": "Tensorboard support has been moved to the extension [Tensorboard extension](https://marketplace.visualstudio.com/items?itemName=ms-toolsai.tensorboard). Instead use the setting `tensorBoard.logDirectory`.",
     "python.tensorBoard.logDirectory.deprecationMessage": "Tensorboard support has been moved to the extension Tensorboard extension. Instead use the setting `tensorBoard.logDirectory`.",
+    "python.terminal.enableSITerminal.description": "Enable Shell Integration for Python Terminal REPL.",
     "python.terminal.activateEnvInCurrentTerminal.description": "Activate Python Environment in the current Terminal on load of the Extension.",
     "python.terminal.activateEnvironment.description": "Activate Python Environment in all Terminals created.",
     "python.terminal.executeInFileDir.description": "When executing a file in the terminal, whether to use execute in the file's directory, instead of the current open folder.",

src/client/common/configSettings.ts

@@ -368,6 +368,7 @@ export class PythonSettings implements IPythonSettings {
                   launchArgs: [],
                   activateEnvironment: true,
                   activateEnvInCurrentTerminal: false,
+                  pythonrcStartup: true,
               };
 
         this.REPL = pythonSettings.get<IREPLSettings>('REPL')!;

src/client/common/terminal/service.ts

@@ -75,6 +75,7 @@ export class TerminalService implements ITerminalService, Disposable {
         }
         this.terminal!.sendText(text);
     }
+
     public async executeCommand(commandLine: string): Promise<ITerminalExecutedCommand | undefined> {
         const terminal = this.terminal!;
         if (!this.options?.hideFromUser) {
@@ -125,6 +126,9 @@ export class TerminalService implements ITerminalService, Disposable {
             this.terminal!.show(preserveFocus);
         }
     }
+    //     When terminal is launched
+    //      Copy pythonrc into <userdata>/pythonrc.py
+    //      Update environment variable collection to include PYTHONSTARTUP=<userdata>/pythonrc.py
     // TODO: Debt switch to Promise<Terminal> ---> breaks 20 tests
     public async ensureTerminal(preserveFocus: boolean = true): Promise<void> {
         if (this.terminal) {

src/client/common/types.ts

@@ -197,6 +197,7 @@ export interface ITerminalSettings {
     readonly launchArgs: string[];
     readonly activateEnvironment: boolean;
     readonly activateEnvInCurrentTerminal: boolean;
+    readonly pythonrcStartup: boolean;
 }
 
 export interface IREPLSettings {

src/client/extensionActivation.ts

@@ -32,7 +32,12 @@ import { TerminalProvider } from './providers/terminalProvider';
 import { setExtensionInstallTelemetryProperties } from './telemetry/extensionInstallTelemetry';
 import { registerTypes as tensorBoardRegisterTypes } from './tensorBoard/serviceRegistry';
 import { registerTypes as commonRegisterTerminalTypes } from './terminals/serviceRegistry';
-import { ICodeExecutionHelper, ICodeExecutionManager, ITerminalAutoActivation } from './terminals/types';
+import {
+    ICodeExecutionHelper,
+    ICodeExecutionManager,
+    IPythonStartupEnvVarService,
+    ITerminalAutoActivation,
+} from './terminals/types';
 import { registerTypes as unitTestsRegisterTypes } from './testing/serviceRegistry';
 
 // components
@@ -177,6 +182,8 @@ async function activateLegacy(ext: ExtensionState, startupStopWatch: StopWatch):
 
             serviceManager.get<ITerminalAutoActivation>(ITerminalAutoActivation).register();
 
+            serviceManager.get<IPythonStartupEnvVarService>(IPythonStartupEnvVarService).register();
+
             serviceManager.get<ICodeExecutionManager>(ICodeExecutionManager).registerCommands();
 
             disposables.push(new ReplProvider(serviceContainer));

src/client/terminals/envCollectionActivation/service.ts

@@ -37,7 +37,11 @@ import { TerminalShellType } from '../../common/terminal/types';
 import { OSType } from '../../common/utils/platform';
 
 import { PythonEnvType } from '../../pythonEnvironments/base/info';
-import { IShellIntegrationService, ITerminalDeactivateService, ITerminalEnvVarCollectionService } from '../types';
+import {
+    IShellIntegrationDetectionService,
+    ITerminalDeactivateService,
+    ITerminalEnvVarCollectionService,
+} from '../types';
 import { ProgressService } from '../../common/application/progressService';
 
 @injectable()
@@ -80,7 +84,8 @@ export class TerminalEnvVarCollectionService implements IExtensionActivationServ
         @inject(IConfigurationService) private readonly configurationService: IConfigurationService,
         @inject(ITerminalDeactivateService) private readonly terminalDeactivateService: ITerminalDeactivateService,
         @inject(IPathUtils) private readonly pathUtils: IPathUtils,
-        @inject(IShellIntegrationService) private readonly shellIntegrationService: IShellIntegrationService,
+        @inject(IShellIntegrationDetectionService)
+        private readonly shellIntegrationService: IShellIntegrationDetectionService,
         @inject(IEnvironmentVariablesProvider)
         private readonly environmentVariablesProvider: IEnvironmentVariablesProvider,
     ) {

src/client/terminals/envCollectionActivation/shellIntegrationService.ts

@@ -14,7 +14,7 @@ import { TerminalShellType } from '../../common/terminal/types';
 import { IDisposableRegistry, IPersistentStateFactory } from '../../common/types';
 import { sleep } from '../../common/utils/async';
 import { traceError, traceVerbose } from '../../logging';
-import { IShellIntegrationService } from '../types';
+import { IShellIntegrationDetectionService } from '../types';
 
 /**
  * This is a list of shells which support shell integration:
@@ -33,7 +33,7 @@ export enum isShellIntegrationWorking {
 }
 
 @injectable()
-export class ShellIntegrationService implements IShellIntegrationService {
+export class ShellIntegrationDetectionService implements IShellIntegrationDetectionService {
     private isWorkingForShell = new Set<TerminalShellType>();
 
     private readonly didChange = new EventEmitter<void>();

src/client/terminals/pythonStartupEnvVar/service.ts

@@ -0,0 +1,43 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import * as path from 'path';
+import { injectable, inject } from 'inversify';
+import { Uri, workspace } from 'vscode';
+import { IPythonStartupEnvVarService } from '../types';
+import { IConfigurationService, IExtensionContext } from '../../common/types';
+import { EXTENSION_ROOT_DIR } from '../../constants';
+
+@injectable()
+export class PythonStartupEnvVarService implements IPythonStartupEnvVarService {
+    constructor(
+        @inject(IExtensionContext) private context: IExtensionContext,
+        @inject(IConfigurationService) private configurationService: IConfigurationService,
+    ) {}
+
+    public async register(): Promise<void> {
+        const storageUri = this.context.storageUri || this.context.globalStorageUri;
+        try {
+            await workspace.fs.createDirectory(storageUri);
+        } catch {
+            // already exists, most likely
+        }
+        const destPath = Uri.joinPath(storageUri, 'pythonrc.py');
+
+        // TODO: Only do this when we have a setting
+        // Rollout strategy:
+        // Stage 1. Opt-in setting in stable/insiders
+        // Stage 2. Out-out setting in insiders
+        // Stage 3. Out-out setting in stable (or experiment?)
+        const sourcePath = path.join(EXTENSION_ROOT_DIR, 'python_files', 'pythonrc.py');
+
+        await workspace.fs.copy(Uri.file(sourcePath), destPath, { overwrite: true });
+        const pythonrcSetting = this.configurationService.getSettings().terminal.pythonrcStartup;
+        // TODO: Is there better place to set this?
+        if (pythonrcSetting) {
+            this.context.environmentVariableCollection.replace('PYTHONSTARTUP', destPath.fsPath);
+        } else {
+            this.context.environmentVariableCollection.delete('PYTHONSTARTUP');
+        }
+    }
+}

src/client/terminals/serviceRegistry.ts

@@ -12,16 +12,18 @@ import {
     ICodeExecutionHelper,
     ICodeExecutionManager,
     ICodeExecutionService,
-    IShellIntegrationService,
+    IPythonStartupEnvVarService,
+    IShellIntegrationDetectionService,
     ITerminalAutoActivation,
     ITerminalDeactivateService,
     ITerminalEnvVarCollectionService,
 } from './types';
 import { TerminalEnvVarCollectionService } from './envCollectionActivation/service';
 import { IExtensionActivationService, IExtensionSingleActivationService } from '../activation/types';
 import { TerminalIndicatorPrompt } from './envCollectionActivation/indicatorPrompt';
-import { ShellIntegrationService } from './envCollectionActivation/shellIntegrationService';
 import { TerminalDeactivateService } from './envCollectionActivation/deactivateService';
+import { ShellIntegrationDetectionService } from './envCollectionActivation/shellIntegrationService';
+import { PythonStartupEnvVarService } from './pythonStartupEnvVar/service';
 
 export function registerTypes(serviceManager: IServiceManager): void {
     serviceManager.addSingleton<ICodeExecutionHelper>(ICodeExecutionHelper, CodeExecutionHelper);
@@ -50,6 +52,11 @@ export function registerTypes(serviceManager: IServiceManager): void {
         IExtensionSingleActivationService,
         TerminalIndicatorPrompt,
     );
-    serviceManager.addSingleton<IShellIntegrationService>(IShellIntegrationService, ShellIntegrationService);
+    serviceManager.addSingleton<IShellIntegrationDetectionService>(
+        IShellIntegrationDetectionService,
+        ShellIntegrationDetectionService,
+    );
+    serviceManager.addSingleton<IPythonStartupEnvVarService>(IPythonStartupEnvVarService, PythonStartupEnvVarService);
+
     serviceManager.addBinding(ITerminalEnvVarCollectionService, IExtensionActivationService);
 }

src/client/terminals/types.ts

@@ -42,8 +42,8 @@ export interface ITerminalEnvVarCollectionService {
     isTerminalPromptSetCorrectly(resource?: Resource): boolean;
 }
 
-export const IShellIntegrationService = Symbol('IShellIntegrationService');
-export interface IShellIntegrationService {
+export const IShellIntegrationDetectionService = Symbol('IShellIntegrationDetectionService');
+export interface IShellIntegrationDetectionService {
     onDidChangeStatus: Event<void>;
     isWorking(): Promise<boolean>;
 }
@@ -53,3 +53,8 @@ export interface ITerminalDeactivateService {
     initializeScriptParams(shell: string): Promise<void>;
     getScriptLocation(shell: string, resource: Resource): Promise<string | undefined>;
 }
+
+export const IPythonStartupEnvVarService = Symbol('IPythonStartupEnvVarService');
+export interface IPythonStartupEnvVarService {
+    register(): void;
+}

src/test/interpreters/activation/terminalEnvVarCollectionService.unit.test.ts

@@ -37,7 +37,7 @@ import { IInterpreterService } from '../../../client/interpreter/contracts';
 import { PathUtils } from '../../../client/common/platform/pathUtils';
 import { PythonEnvType } from '../../../client/pythonEnvironments/base/info';
 import { PythonEnvironment } from '../../../client/pythonEnvironments/info';
-import { IShellIntegrationService, ITerminalDeactivateService } from '../../../client/terminals/types';
+import { IShellIntegrationDetectionService, ITerminalDeactivateService } from '../../../client/terminals/types';
 import { IEnvironmentVariablesProvider } from '../../../client/common/variables/types';
 
 suite('Terminal Environment Variable Collection Service', () => {
@@ -58,7 +58,7 @@ suite('Terminal Environment Variable Collection Service', () => {
         title: Interpreters.activatingTerminals,
     };
     let configService: IConfigurationService;
-    let shellIntegrationService: IShellIntegrationService;
+    let shellIntegrationService: IShellIntegrationDetectionService;
     const displayPath = 'display/path';
     const customShell = 'powershell';
     const defaultShell = defaultShells[getOSType()];
@@ -76,7 +76,7 @@ suite('Terminal Environment Variable Collection Service', () => {
         context = mock<IExtensionContext>();
         shell = mock<IApplicationShell>();
         const envVarProvider = mock<IEnvironmentVariablesProvider>();
-        shellIntegrationService = mock<IShellIntegrationService>();
+        shellIntegrationService = mock<IShellIntegrationDetectionService>();
         when(shellIntegrationService.isWorking()).thenResolve(true);
         globalCollection = mock<GlobalEnvironmentVariableCollection>();
         collection = mock<EnvironmentVariableCollection>();

src/test/terminals/serviceRegistry.unit.test.ts

@@ -14,17 +14,20 @@ import { TerminalCodeExecutionProvider } from '../../client/terminals/codeExecut
 import { TerminalDeactivateService } from '../../client/terminals/envCollectionActivation/deactivateService';
 import { TerminalIndicatorPrompt } from '../../client/terminals/envCollectionActivation/indicatorPrompt';
 import { TerminalEnvVarCollectionService } from '../../client/terminals/envCollectionActivation/service';
-import { ShellIntegrationService } from '../../client/terminals/envCollectionActivation/shellIntegrationService';
+
 import { registerTypes } from '../../client/terminals/serviceRegistry';
 import {
     ICodeExecutionHelper,
     ICodeExecutionManager,
     ICodeExecutionService,
-    IShellIntegrationService,
+    IPythonStartupEnvVarService,
+    IShellIntegrationDetectionService,
     ITerminalAutoActivation,
     ITerminalDeactivateService,
     ITerminalEnvVarCollectionService,
 } from '../../client/terminals/types';
+import { ShellIntegrationDetectionService } from '../../client/terminals/envCollectionActivation/shellIntegrationService';
+import { PythonStartupEnvVarService } from '../../client/terminals/pythonStartupEnvVar/service';
 
 suite('Terminal - Service Registry', () => {
     test('Ensure all services get registered', () => {
@@ -39,7 +42,8 @@ suite('Terminal - Service Registry', () => {
             [ITerminalEnvVarCollectionService, TerminalEnvVarCollectionService],
             [IExtensionSingleActivationService, TerminalIndicatorPrompt],
             [ITerminalDeactivateService, TerminalDeactivateService],
-            [IShellIntegrationService, ShellIntegrationService],
+            [IShellIntegrationDetectionService, ShellIntegrationDetectionService],
+            [IPythonStartupEnvVarService, PythonStartupEnvVarService],
         ].forEach((args) => {
             if (args.length === 2) {
                 services