Add Python activate hooks into shell integration scripts (microsoft#261094)

* Add activate hook for bash

* Add bash Python activate hooks

* Register zsh activate hooks

* Add hooks for .ps1

* clean up for pwsh

* Add /^VSCODE_(PWSH|ZSH|BASH)_ACTIVATE$/,

* Do not let non-Python-env ext to set activate var

* Remove comment

* Add logService --> Need review!

* Add fish activate hooks

* enable log in envVarCollection.ts

* Attach PYTHON after VSCODE for activate hook name

* make blockPythonActivationVar func

* Revert logService

* Show the * next to error message.

* clean-up

* Dont forget to rename

* Dont mess with original formatting

* Still need .+, remove dup Python keysToRemove

* Add tests for Python activate vars

* Consolidate tests!

* Clean clean clean

Author: anthonykim1

extensions/terminal-suggest/src/terminalSuggestMain.ts

@@ -520,7 +520,7 @@ export function sanitizeProcessEnvironment(env: Record<string, string>, ...prese
 	}, {});
 	const keysToRemove = [
 		/^ELECTRON_.$/,
-		/^VSCODE_(?!(PORTABLE|SHELL_LOGIN|ENV_REPLACE|ENV_APPEND|ENV_PREPEND)).$/,
+		/^VSCODE_(?!(PORTABLE|SHELL_LOGIN|ENV_REPLACE|ENV_APPEND|ENV_PREPEND)).+$/,
 		/^SNAP(|_.*)$/,
 		/^GDK_PIXBUF_.$/,
 	];

src/vs/base/test/common/processes.test.ts

@@ -28,12 +28,20 @@ suite('Processes', () => {
 			VSCODE_CODE_CACHE_PATH: 'x',
 			VSCODE_NEW_VAR: 'x',
 			GDK_PIXBUF_MODULE_FILE: 'x',
-			GDK_PIXBUF_MODULEDIR: 'x'
+			GDK_PIXBUF_MODULEDIR: 'x',
+			VSCODE_PYTHON_BASH_ACTIVATE: 'source /path/to/venv/bin/activate',
+			VSCODE_PYTHON_ZSH_ACTIVATE: 'source /path/to/venv/bin/activate',
+			VSCODE_PYTHON_PWSH_ACTIVATE: '. /path/to/venv/Scripts/Activate.ps1',
+			VSCODE_PYTHON_FISH_ACTIVATE: 'source /path/to/venv/bin/activate.fish',
 		};
 		processes.sanitizeProcessEnvironment(env);
 		assert.strictEqual(env['FOO'], 'bar');
 		assert.strictEqual(env['VSCODE_SHELL_LOGIN'], '1');
 		assert.strictEqual(env['VSCODE_PORTABLE'], '3');
+		assert.strictEqual(env['VSCODE_PYTHON_BASH_ACTIVATE'], undefined);
+		assert.strictEqual(env['VSCODE_PYTHON_ZSH_ACTIVATE'], undefined);
+		assert.strictEqual(env['VSCODE_PYTHON_PWSH_ACTIVATE'], undefined);
+		assert.strictEqual(env['VSCODE_PYTHON_FISH_ACTIVATE'], undefined);
 		assert.strictEqual(Object.keys(env).length, 3);
 	});
 });

src/vs/platform/terminal/common/environmentVariableCollection.ts

@@ -13,6 +13,8 @@ const mutatorTypeToLabelMap: Map<EnvironmentVariableMutatorType, string> = new M
 	[EnvironmentVariableMutatorType.Prepend, 'PREPEND'],
 	[EnvironmentVariableMutatorType.Replace, 'REPLACE']
 ]);
+const PYTHON_ACTIVATION_VARS_PATTERN = /^VSCODE_PYTHON_(PWSH|ZSH|BASH|FISH)_ACTIVATE/;
+const PYTHON_ENV_EXTENSION_ID = 'ms-python.vscode-python-envs';
 
 export class MergedEnvironmentVariableCollection implements IMergedEnvironmentVariableCollection {
 	private readonly map: Map<string, IExtensionOwnedEnvironmentVariableMutator[]> = new Map();
@@ -28,6 +30,12 @@ export class MergedEnvironmentVariableCollection implements IMergedEnvironmentVa
 			while (!next.done) {
 				const mutator = next.value[1];
 				const key = next.value[0];
+
+				if (this.blockPythonActivationVar(key, extensionIdentifier)) {
+					next = it.next();
+					continue;
+				}
+
 				let entry = this.map.get(key);
 				if (!entry) {
 					entry = [];
@@ -70,6 +78,11 @@ export class MergedEnvironmentVariableCollection implements IMergedEnvironmentVa
 			const actualVariable = isWindows ? lowerToActualVariableNames![variable.toLowerCase()] || variable : variable;
 			for (const mutator of mutators) {
 				const value = variableResolver ? await variableResolver(mutator.value) : mutator.value;
+
+				if (this.blockPythonActivationVar(mutator.variable, mutator.extensionIdentifier)) {
+					continue;
+				}
+
 				// Default: true
 				if (mutator.options?.applyAtProcessCreation ?? true) {
 					switch (mutator.type) {
@@ -97,6 +110,14 @@ export class MergedEnvironmentVariableCollection implements IMergedEnvironmentVa
 		return value.replaceAll(':', '\\x3a');
 	}
 
+	private blockPythonActivationVar(variable: string, extensionIdentifier: string): boolean {
+		// Only Python env extension can modify Python activate env var.
+		if (PYTHON_ACTIVATION_VARS_PATTERN.test(variable) && PYTHON_ENV_EXTENSION_ID !== extensionIdentifier) {
+			return true;
+		}
+		return false;
+	}
+
 	diff(other: IMergedEnvironmentVariableCollection, scope: EnvironmentVariableScope | undefined): IMergedEnvironmentVariableCollectionDiff | undefined {
 		const added: Map<string, IExtensionOwnedEnvironmentVariableMutator[]> = new Map();
 		const changed: Map<string, IExtensionOwnedEnvironmentVariableMutator[]> = new Map();

src/vs/workbench/contrib/terminal/common/scripts/shellIntegration-bash.sh

@@ -91,6 +91,15 @@ if [ -n "${VSCODE_ENV_APPEND:-}" ]; then
 	builtin unset VSCODE_ENV_APPEND
 fi
 
+# Register Python shell activate hooks
+if [ -n "$VSCODE_PYTHON_BASH_ACTIVATE" ] && [ "$TERM_PROGRAM" = "vscode" ]; then
+	# Prevent crashing by negating exit code
+	if ! builtin eval "$VSCODE_PYTHON_BASH_ACTIVATE"; then
+		__vsc_activation_status=$?
+		builtin printf '\x1b[0m\x1b[7m * \x1b[0;103m VS Code Python bash activation failed with exit code %d \x1b[0m' "$__vsc_activation_status"
+	fi
+fi
+
 __vsc_get_trap() {
 	# 'trap -p DEBUG' outputs a shell command like `trap -- '…shellcode…' DEBUG`.
 	# The terms are quoted literals, but are not guaranteed to be on a single line.

src/vs/workbench/contrib/terminal/common/scripts/shellIntegration-rc.zsh

@@ -68,6 +68,15 @@ if [ -n "${VSCODE_ENV_APPEND:-}" ]; then
 	unset VSCODE_ENV_APPEND
 fi
 
+# Register Python shell activate hooks
+if [ -n "$VSCODE_PYTHON_ZSH_ACTIVATE" ] && [ "$TERM_PROGRAM" = "vscode" ]; then
+	# Prevent crashing by negating exit code
+	if ! builtin eval "$VSCODE_PYTHON_ZSH_ACTIVATE"; then
+		__vsc_activation_status=$?
+		builtin printf '\x1b[0m\x1b[7m * \x1b[0;103m VS Code Python zsh activation failed with exit code %d \x1b[0m' "$__vsc_activation_status"
+	fi
+fi
+
 # Report prompt type
 if [ -n "$ZSH" ] && [ -n "$ZSH_VERSION" ] && (( ${+functions[omz]} )) ; then
 	builtin printf '\e]633;P;PromptType=oh-my-zsh\a'

src/vs/workbench/contrib/terminal/common/scripts/shellIntegration.fish

@@ -75,6 +75,17 @@ function __vsc_apply_env_vars
 	end
 end
 
+# Register Python shell activate hooks
+if test -n "$VSCODE_PYTHON_FISH_ACTIVATE"; and test "$TERM_PROGRAM" = "vscode"
+	# Fish does not crash on eval failure, so don't need negation.
+	eval $VSCODE_PYTHON_FISH_ACTIVATE
+	set __vsc_activation_status $status
+
+	if test $__vsc_activation_status -ne 0
+		builtin printf '\x1b[0m\x1b[7m * \x1b[0;103m VS Code Python fish activation failed with exit code %d \x1b[0m \n' "$__vsc_activation_status"
+	end
+end
+
 # Handle the shell integration nonce
 if set -q VSCODE_NONCE
 	set -l __vsc_nonce $VSCODE_NONCE

src/vs/workbench/contrib/terminal/common/scripts/shellIntegration.ps1

@@ -68,6 +68,20 @@ if ($env:VSCODE_ENV_APPEND) {
 	$env:VSCODE_ENV_APPEND = $null
 }
 
+# Register Python shell activate hooks
+if ($env:VSCODE_PYTHON_PWSH_ACTIVATE -and $env:TERM_PROGRAM -eq 'vscode') {
+	$activateScript = $env:VSCODE_PYTHON_PWSH_ACTIVATE
+	Remove-Item Env:VSCODE_PYTHON_PWSH_ACTIVATE
+
+	try {
+		Invoke-Expression $activateScript
+	}
+	catch {
+		$activationError = $_
+		Write-Host "`e[0m`e[7m * `e[0;103m VS Code Python powershell activation failed with exit code $($activationError.Exception.Message) `e[0m"
+	}
+}
+
 function Global:__VSCode-Escape-Value([string]$value) {
 	# NOTE: In PowerShell v6.1+, this can be written `$value -replace '…', { … }` instead of `[regex]::Replace`.
 	# Replace any non-alphanumeric characters.