Environment Variable management and bug fixes (#28)

Author: karthiknadig

package-lock.json

@@ -451,6 +451,7 @@
         "@iarna/toml": "^2.2.5",
         "@vscode/extension-telemetry": "^0.9.7",
         "@vscode/test-cli": "^0.0.10",
+        "dotenv": "^16.4.5",
         "fs-extra": "^11.2.0",
         "stack-trace": "0.0.10",
         "vscode-jsonrpc": "^9.0.0-next.5",

package.json

@@ -8,6 +8,7 @@ import {
     Terminal,
     TaskExecution,
     TerminalOptions,
+    FileChangeType,
 } from 'vscode';
 
 /**
@@ -1105,11 +1106,45 @@ export interface PythonExecutionApi
         PythonTerminalRunApi,
         PythonTaskRunApi,
         PythonBackgroundRunApi {}
+
+export interface DidChangeEnvironmentVariablesEventArgs {
+    uri?: Uri;
+    changeTye: FileChangeType;
+}
+
+export interface PythonEnvironmentVariablesApi {
+    /**
+     * Get environment variables for a workspace. This picks up `.env` file from the root of the
+     * workspace.
+     *
+     * Order of overrides:
+     * 1. `baseEnvVar` if given or `process.env`
+     * 2. `.env` file from the "python.envFile" setting in the workspace.
+     * 3. `.env` file at the root of the python project.
+     * 4. `overrides` in the order provided.
+     *
+     * @param uri The URI of the project, workspace or a file in a for which environment variables are required.
+     * @param overrides Additional environment variables to override the defaults.
+     * @param baseEnvVar The base environment variables that should be used as a starting point.
+     */
+    getEnvironmentVariables(
+        uri: Uri,
+        overrides?: ({ [key: string]: string | undefined } | Uri)[],
+        baseEnvVar?: { [key: string]: string | undefined },
+    ): Promise<{ [key: string]: string | undefined }>;
+
+    /**
+     * Event raised when `.env` file changes or any other monitored source of env variable changes.
+     */
+    onDidChangeEnvironmentVariables: Event<DidChangeEnvironmentVariablesEventArgs>;
+}
+
 /**
  * The API for interacting with Python environments, package managers, and projects.
  */
 export interface PythonEnvironmentApi
     extends PythonEnvironmentManagerApi,
         PythonPackageManagerApi,
         PythonProjectApi,
-        PythonExecutionApi {}
+        PythonExecutionApi,
+        PythonEnvironmentVariablesApi {}

src/api.ts

@@ -0,0 +1,40 @@
+import { Uri } from 'vscode';
+import { getWorkspaceFolder, getWorkspaceFolders } from '../workspace.apis';
+
+export function resolveVariables(value: string, project?: Uri, env?: { [key: string]: string }): string {
+    const substitutions = new Map<string, string>();
+    const home = process.env.HOME || process.env.USERPROFILE;
+    if (home) {
+        substitutions.set('${userHome}', home);
+    }
+
+    if (project) {
+        substitutions.set('${pythonProject}', project.fsPath);
+    }
+
+    const workspace = project ? getWorkspaceFolder(project) : undefined;
+    if (workspace) {
+        substitutions.set('${workspaceFolder}', workspace.uri.fsPath);
+    }
+    substitutions.set('${cwd}', process.cwd());
+    (getWorkspaceFolders() ?? []).forEach((w) => {
+        substitutions.set('${workspaceFolder:' + w.name + '}', w.uri.fsPath);
+    });
+
+    const substEnv = env || process.env;
+    if (substEnv) {
+        for (const [key, value] of Object.entries(substEnv)) {
+            if (value && key.length > 0) {
+                substitutions.set('${env:' + key + '}', value);
+            }
+        }
+    }
+
+    let result = value;
+    substitutions.forEach((v, k) => {
+        while (k.length > 0 && result.indexOf(k) >= 0) {
+            result = result.replace(k, v);
+        }
+    });
+    return result;
+}

src/common/utils/internalVariables.ts

@@ -0,0 +1,9 @@
+import { FileStat, Uri, workspace } from 'vscode';
+
+export function readFile(uri: Uri): Thenable<Uint8Array> {
+    return workspace.fs.readFile(uri);
+}
+
+export function stat(uri: Uri): Thenable<FileStat> {
+    return workspace.fs.stat(uri);
+}

src/common/workspace.fs.apis.ts

@@ -51,6 +51,7 @@ import {
 } from './features/terminal/activateMenuButton';
 import { PythonStatusBarImpl } from './features/views/pythonStatusBar';
 import { updateViewsAndStatus } from './features/views/revealHandler';
+import { EnvVarManager, PythonEnvVariableManager } from './features/execution/envVariableManager';
 
 export async function activate(context: ExtensionContext): Promise<PythonEnvironmentApi> {
     // Logging should be set up before anything else.
@@ -69,6 +70,9 @@ export async function activate(context: ExtensionContext): Promise<PythonEnviron
     const projectManager: PythonProjectManager = new PythonProjectManagerImpl();
     context.subscriptions.push(projectManager);
 
+    const envVarManager: EnvVarManager = new PythonEnvVariableManager(projectManager);
+    context.subscriptions.push(envVarManager);
+
     const envManagers: EnvironmentManagers = new PythonEnvironmentManagers(projectManager);
     context.subscriptions.push(envManagers);
 
@@ -79,7 +83,7 @@ export async function activate(context: ExtensionContext): Promise<PythonEnviron
         registerAutoProjectProvider(projectCreators),
     );
 
-    setPythonApi(envManagers, projectManager, projectCreators, terminalManager);
+    setPythonApi(envManagers, projectManager, projectCreators, terminalManager, envVarManager);
 
     const managerView = new EnvManagerView(envManagers);
     context.subscriptions.push(managerView);
@@ -103,10 +107,10 @@ export async function activate(context: ExtensionContext): Promise<PythonEnviron
             await refreshPackagesCommand(item);
         }),
         commands.registerCommand('python-envs.create', async (item) => {
-            await createEnvironmentCommand(item, envManagers, projectManager);
+            return await createEnvironmentCommand(item, envManagers, projectManager);
         }),
         commands.registerCommand('python-envs.createAny', async () => {
-            await createAnyEnvironmentCommand(envManagers, projectManager);
+            return await createAnyEnvironmentCommand(envManagers, projectManager);
         }),
         commands.registerCommand('python-envs.remove', async (item) => {
             await removeEnvironmentCommand(item, envManagers);

src/extension.ts

@@ -68,48 +68,58 @@ export async function createEnvironmentCommand(
     context: unknown,
     em: EnvironmentManagers,
     pm: PythonProjectManager,
-): Promise<void> {
+): Promise<PythonEnvironment | undefined> {
     if (context instanceof EnvManagerTreeItem) {
         const manager = (context as EnvManagerTreeItem).manager;
         const projects = await pickProjectMany(pm.getProjects());
         if (projects) {
-            await manager.create(projects.length === 0 ? 'global' : projects.map((p) => p.uri));
+            return await manager.create(projects.length === 0 ? 'global' : projects.map((p) => p.uri));
+        } else {
+            traceError(`No projects found for ${context}`);
         }
     } else if (context instanceof Uri) {
         const manager = em.getEnvironmentManager(context as Uri);
         const project = pm.get(context as Uri);
         if (project) {
-            await manager?.create(project.uri);
+            return await manager?.create(project.uri);
+        } else {
+            traceError(`No project found for ${context}`);
         }
     } else {
         traceError(`Invalid context for create command: ${context}`);
     }
 }
 
-export async function createAnyEnvironmentCommand(em: EnvironmentManagers, pm: PythonProjectManager): Promise<void> {
+export async function createAnyEnvironmentCommand(
+    em: EnvironmentManagers,
+    pm: PythonProjectManager,
+): Promise<PythonEnvironment | undefined> {
     const projects = await pickProjectMany(pm.getProjects());
     if (projects && projects.length > 0) {
         const defaultManagers: InternalEnvironmentManager[] = [];
+
         projects.forEach((p) => {
             const manager = em.getEnvironmentManager(p.uri);
             if (manager && manager.supportsCreate && !defaultManagers.includes(manager)) {
                 defaultManagers.push(manager);
             }
         });
+
         const managerId = await pickEnvironmentManager(
             em.managers.filter((m) => m.supportsCreate),
             defaultManagers,
         );
 
         const manager = em.managers.find((m) => m.id === managerId);
         if (manager) {
-            await manager.create(projects.map((p) => p.uri));
+            return await manager.create(projects.map((p) => p.uri));
         }
     } else if (projects && projects.length === 0) {
         const managerId = await pickEnvironmentManager(em.managers.filter((m) => m.supportsCreate));
+
         const manager = em.managers.find((m) => m.id === managerId);
         if (manager) {
-            await manager.create('global');
+            return await manager.create('global');
         }
     }
 }

src/features/envCommands.ts

@@ -0,0 +1,34 @@
+import { Uri } from 'vscode';
+import { readFile } from '../../common/workspace.fs.apis';
+import { parse } from 'dotenv';
+
+export function mergeEnvVariables(
+    base: { [key: string]: string | undefined },
+    other: { [key: string]: string | undefined },
+) {
+    const env: { [key: string]: string | undefined } = {};
+
+    Object.keys(other).forEach((otherKey) => {
+        let value = other[otherKey];
+        if (value === undefined || value === '') {
+            // SOME_ENV_VAR=
+            delete env[otherKey];
+        } else {
+            Object.keys(base).forEach((baseKey) => {
+                const baseValue = base[baseKey];
+                if (baseValue) {
+                    value = value?.replace(`\${${baseKey}}`, baseValue);
+                }
+            });
+            env[otherKey] = value;
+        }
+    });
+
+    return env;
+}
+
+export async function parseEnvFile(envFile: Uri): Promise<{ [key: string]: string | undefined }> {
+    const raw = await readFile(envFile);
+    const contents = Buffer.from(raw).toString('utf-8');
+    return parse(contents);
+}

src/features/execution/envVarUtils.ts

@@ -0,0 +1,83 @@
+import * as path from 'path';
+import * as fsapi from 'fs-extra';
+import { Uri, Event, EventEmitter, FileChangeType } from 'vscode';
+import { DidChangeEnvironmentVariablesEventArgs, PythonEnvironmentVariablesApi } from '../../api';
+import { Disposable } from 'vscode-jsonrpc';
+import { createFileSystemWatcher, getConfiguration } from '../../common/workspace.apis';
+import { PythonProjectManager } from '../../internal.api';
+import { mergeEnvVariables, parseEnvFile } from './envVarUtils';
+import { resolveVariables } from '../../common/utils/internalVariables';
+
+export interface EnvVarManager extends PythonEnvironmentVariablesApi, Disposable {}
+
+export class PythonEnvVariableManager implements EnvVarManager {
+    private disposables: Disposable[] = [];
+
+    private _onDidChangeEnvironmentVariables;
+    private watcher;
+
+    constructor(private pm: PythonProjectManager) {
+        this._onDidChangeEnvironmentVariables = new EventEmitter<DidChangeEnvironmentVariablesEventArgs>();
+        this.onDidChangeEnvironmentVariables = this._onDidChangeEnvironmentVariables.event;
+
+        this.watcher = createFileSystemWatcher('**/.env');
+        this.disposables.push(
+            this._onDidChangeEnvironmentVariables,
+            this.watcher,
+            this.watcher.onDidCreate((e) =>
+                this._onDidChangeEnvironmentVariables.fire({ uri: e, changeTye: FileChangeType.Created }),
+            ),
+            this.watcher.onDidChange((e) =>
+                this._onDidChangeEnvironmentVariables.fire({ uri: e, changeTye: FileChangeType.Changed }),
+            ),
+            this.watcher.onDidDelete((e) =>
+                this._onDidChangeEnvironmentVariables.fire({ uri: e, changeTye: FileChangeType.Deleted }),
+            ),
+        );
+    }
+
+    async getEnvironmentVariables(
+        uri: Uri,
+        overrides?: ({ [key: string]: string | undefined } | Uri)[],
+        baseEnvVar?: { [key: string]: string | undefined },
+    ): Promise<{ [key: string]: string | undefined }> {
+        const project = this.pm.get(uri);
+
+        const base = baseEnvVar || { ...process.env };
+        let env = base;
+
+        const config = getConfiguration('python', project?.uri ?? uri);
+        let envFilePath = config.get<string>('envFile');
+        envFilePath = envFilePath ? path.normalize(resolveVariables(envFilePath)) : undefined;
+
+        if (envFilePath && (await fsapi.pathExists(envFilePath))) {
+            const other = await parseEnvFile(Uri.file(envFilePath));
+            env = mergeEnvVariables(env, other);
+        }
+
+        let projectEnvFilePath = project ? path.normalize(path.join(project.uri.fsPath, '.env')) : undefined;
+        if (
+            projectEnvFilePath &&
+            projectEnvFilePath?.toLowerCase() !== envFilePath?.toLowerCase() &&
+            (await fsapi.pathExists(projectEnvFilePath))
+        ) {
+            const other = await parseEnvFile(Uri.file(projectEnvFilePath));
+            env = mergeEnvVariables(env, other);
+        }
+
+        if (overrides) {
+            for (const override of overrides) {
+                const other = override instanceof Uri ? await parseEnvFile(override) : override;
+                env = mergeEnvVariables(env, other);
+            }
+        }
+
+        return env;
+    }
+
+    onDidChangeEnvironmentVariables: Event<DidChangeEnvironmentVariablesEventArgs>;
+
+    dispose(): void {
+        this.disposables.forEach((disposable) => disposable.dispose());
+    }
+}