From fb4fccf2d5fa5086ce53e22f83ccdc7dee81a13d Mon Sep 17 00:00:00 2001
From: Chris LaFreniere <40371649+chlafreniere@users.noreply.github.com>
Date: Wed, 23 Oct 2019 19:22:24 -0700
Subject: [PATCH] Notebooks: ensure python path dirs added to path on session
 start (#7968)

* ensure python path dirs add to path session start

* Change logic slightly

* PR feedback from Charles
---
 .../notebook/src/jupyter/jupyterNotebookManager.ts |  3 ++-
 .../notebook/src/jupyter/jupyterServerManager.ts   |  4 ++++
 .../notebook/src/jupyter/jupyterSessionManager.ts  | 14 +++++++++-----
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/extensions/notebook/src/jupyter/jupyterNotebookManager.ts b/extensions/notebook/src/jupyter/jupyterNotebookManager.ts
index a12fbedacc..58cdefb0d8 100644
--- a/extensions/notebook/src/jupyter/jupyterNotebookManager.ts
+++ b/extensions/notebook/src/jupyter/jupyterNotebookManager.ts
@@ -18,7 +18,8 @@ export class JupyterNotebookManager implements nb.NotebookManager, vscode.Dispos
 	private _sessionManager: JupyterSessionManager;
 
 	constructor(private _serverManager: LocalJupyterServerManager, sessionManager?: JupyterSessionManager, private apiWrapper: ApiWrapper = new ApiWrapper()) {
-		this._sessionManager = sessionManager || new JupyterSessionManager();
+		let pythonEnvVarPath = this._serverManager && this._serverManager.jupyterServerInstallation && this._serverManager.jupyterServerInstallation.pythonEnvVarPath;
+		this._sessionManager = sessionManager || new JupyterSessionManager(pythonEnvVarPath);
 		this._serverManager.onServerStarted(() => {
 			this.setServerSettings(this._serverManager.serverSettings);
 		});
diff --git a/extensions/notebook/src/jupyter/jupyterServerManager.ts b/extensions/notebook/src/jupyter/jupyterServerManager.ts
index 5955531d03..0b4d715128 100644
--- a/extensions/notebook/src/jupyter/jupyterServerManager.ts
+++ b/extensions/notebook/src/jupyter/jupyterServerManager.ts
@@ -54,6 +54,10 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo
 		return this._onServerStarted.event;
 	}
 
+	public get jupyterServerInstallation(): JupyterServerInstallation | undefined {
+		return this.options && this.options.jupyterInstallation;
+	}
+
 	public async startServer(): Promise<void> {
 		try {
 			if (!this._jupyterServer) {
diff --git a/extensions/notebook/src/jupyter/jupyterSessionManager.ts b/extensions/notebook/src/jupyter/jupyterSessionManager.ts
index 0a7cf04986..c001579013 100644
--- a/extensions/notebook/src/jupyter/jupyterSessionManager.ts
+++ b/extensions/notebook/src/jupyter/jupyterSessionManager.ts
@@ -69,7 +69,7 @@ export class JupyterSessionManager implements nb.SessionManager {
 	private _sessionManager: Session.IManager;
 	private static _sessions: JupyterSession[] = [];
 
-	constructor() {
+	constructor(private _pythonEnvVarPath?: string) {
 		this._isReady = false;
 		this._ready = new Deferred<void>();
 	}
@@ -123,7 +123,7 @@ export class JupyterSessionManager implements nb.SessionManager {
 			return Promise.reject(new Error(localize('errorStartBeforeReady', "Cannot start a session, the manager is not yet initialized")));
 		}
 		let sessionImpl = await this._sessionManager.startNew(options);
-		let jupyterSession = new JupyterSession(sessionImpl, skipSettingEnvironmentVars);
+		let jupyterSession = new JupyterSession(sessionImpl, skipSettingEnvironmentVars, this._pythonEnvVarPath);
 		await jupyterSession.messagesComplete;
 		let index = JupyterSessionManager._sessions.findIndex(session => session.path === options.path);
 		if (index > -1) {
@@ -170,7 +170,7 @@ export class JupyterSession implements nb.ISession {
 	private _kernel: nb.IKernel;
 	private _messagesComplete: Deferred<void> = new Deferred<void>();
 
-	constructor(private sessionImpl: Session.ISession, skipSettingEnvironmentVars?: boolean) {
+	constructor(private sessionImpl: Session.ISession, skipSettingEnvironmentVars?: boolean, private _pythonEnvVarPath?: string) {
 		this.setEnvironmentVars(skipSettingEnvironmentVars);
 	}
 
@@ -335,8 +335,12 @@ export class JupyterSession implements nb.ISession {
 			let allCode: string = '';
 			for (let i = 0; i < Object.keys(process.env).length; i++) {
 				let key = Object.keys(process.env)[i];
-				// Jupyter doesn't seem to alow for setting multiple variables at once, so doing it with multiple commands
-				allCode += `%set_env ${key}=${process.env[key]}${EOL}`;
+				if (key.toLowerCase() === 'path' && this._pythonEnvVarPath) {
+					allCode += `%set_env ${key}=${this._pythonEnvVarPath}${EOL}`;
+				} else {
+					// Jupyter doesn't seem to alow for setting multiple variables at once, so doing it with multiple commands
+					allCode += `%set_env ${key}=${process.env[key]}${EOL}`;
+				}
 			}
 			let future = this.sessionImpl.kernel.requestExecute({
 				code: allCode