Skip to content

Commit 78e999b

Browse files
mdesmetmdesmet
committed
feat: add Python environment management for VSCode extension
1 parent a677fa0 commit 78e999b

File tree

2 files changed

+297
-0
lines changed

2 files changed

+297
-0
lines changed

‎src/manifest/pythonEnvironment.ts

Lines changed: 244 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,244 @@
1+
import { Disposable, Event, extensions, Uri, workspace } from "vscode";
2+
import { EnvironmentVariables } from "../dbt_integration/domain";
3+
import { provideSingleton } from "../utils";
4+
import { DBTTerminal } from "../dbt_integration/terminal";
5+
import { inject } from "inversify";
6+
7+
type EnvFrom = "process" | "integrated" | "dotenv";
8+
interface PythonExecutionDetails {
9+
getPythonPath: () => string;
10+
onDidChangeExecutionDetails: Event<Uri | undefined>;
11+
getEnvVars: () => EnvironmentVariables;
12+
}
13+
14+
@provideSingleton(PythonEnvironment)
15+
export class PythonEnvironment {
16+
private executionDetails?: PythonExecutionDetails;
17+
private disposables: Disposable[] = [];
18+
private environmentVariableSource: Record<string, EnvFrom> = {};
19+
public allPythonPaths: { path: string; pathType: string }[] = [];
20+
public isPython3: boolean = true;
21+
22+
constructor(
23+
@inject("DBTTerminal")
24+
private dbtTerminal: DBTTerminal,
25+
) {}
26+
27+
async dispose(): Promise<void> {
28+
while (this.disposables.length) {
29+
const x = this.disposables.pop();
30+
if (x) {
31+
x.dispose();
32+
}
33+
}
34+
}
35+
36+
async printEnvVars() {
37+
await this.dbtTerminal.show(true);
38+
const envVars = this.environmentVariables;
39+
this.dbtTerminal.log("Printing environment variables...\r\n");
40+
for (const key in envVars) {
41+
this.dbtTerminal.log(
42+
`${key}=${envVars[key]}\t\tsource:${this.environmentVariableSource[key]}\r\n`,
43+
);
44+
}
45+
}
46+
47+
public get pythonPath() {
48+
return (
49+
this.getResolvedConfigValue("dbtPythonPathOverride") ||
50+
this.executionDetails!.getPythonPath()
51+
);
52+
}
53+
54+
public get environmentVariables(): EnvironmentVariables {
55+
if (!this.executionDetails) {
56+
throw new Error(
57+
"executionDetails is undefined, cannot retrieve environment variables",
58+
);
59+
}
60+
return this.executionDetails.getEnvVars();
61+
}
62+
63+
public get onPythonEnvironmentChanged() {
64+
return this.executionDetails!.onDidChangeExecutionDetails;
65+
}
66+
67+
async initialize(): Promise<void> {
68+
if (this.executionDetails !== undefined) {
69+
return;
70+
}
71+
72+
this.executionDetails = await this.activatePythonExtension();
73+
}
74+
75+
getResolvedConfigValue(key: string) {
76+
const value = workspace.getConfiguration("dbt").get<string>(key, "");
77+
return this.substituteSettingsVariables(value, this.environmentVariables);
78+
}
79+
80+
private substituteSettingsVariables(
81+
value: any,
82+
vsCodeEnv: EnvironmentVariables,
83+
): any {
84+
if (!value) {
85+
return value;
86+
}
87+
if (typeof value !== "string") {
88+
return value;
89+
}
90+
const regexVsCodeEnv = /\$\{env\:(.*?)\}/gm;
91+
let matchResult;
92+
while ((matchResult = regexVsCodeEnv.exec(value)) !== null) {
93+
// This is necessary to avoid infinite loops with zero-width matches
94+
if (matchResult.index === regexVsCodeEnv.lastIndex) {
95+
regexVsCodeEnv.lastIndex++;
96+
}
97+
if (vsCodeEnv[matchResult[1]] !== undefined) {
98+
value = value.replace(
99+
new RegExp(`\\\$\\\{env\\\:${matchResult[1]}\\\}`, "gm"),
100+
vsCodeEnv[matchResult[1]]!,
101+
);
102+
this.dbtTerminal.debug(
103+
"pythonEnvironment:substituteSettingsVariables",
104+
`Picking env var ${matchResult[1]} from ${
105+
this.environmentVariableSource[matchResult[1]]
106+
}`,
107+
);
108+
}
109+
}
110+
value = value.replace(
111+
"${workspaceFolder}",
112+
workspace.workspaceFolders![0].uri.fsPath,
113+
);
114+
return value;
115+
}
116+
117+
private async activatePythonExtension(): Promise<PythonExecutionDetails> {
118+
const extension = extensions.getExtension("ms-python.python")!;
119+
120+
if (!extension.isActive) {
121+
await extension.activate();
122+
}
123+
await extension.exports.ready;
124+
125+
const api = extension.exports;
126+
this.allPythonPaths = await api.environment.getEnvironmentPaths();
127+
const pythonPath = api.settings.getExecutionDetails(workspace.workspaceFile)
128+
.execCommand[0];
129+
const envDetails = await api.environment.getEnvironmentDetails(pythonPath);
130+
this.isPython3 = envDetails?.version[0] === "3";
131+
132+
const dbtInstalledPythonPath: string[] = [];
133+
// TODO: support multiple workspacefolders for python detection
134+
// for (const workspaceFolder of workspace.workspaceFolders || []) {
135+
// const candidatePythonPath = api.settings.getExecutionDetails(
136+
// workspaceFolder.uri,
137+
// ).execCommand[0];
138+
139+
// const dbtInstalledCommand =
140+
// this.dbtCommandFactory.createVerifyDbtInstalledCommand();
141+
// const checkDBTInstalledProcess =
142+
// this.commandProcessExecutionFactory.createCommandProcessExecution({
143+
// command: candidatePythonPath,
144+
// args: dbtInstalledCommand.processExecutionParams.args,
145+
// });
146+
147+
// try {
148+
// await checkDBTInstalledProcess.complete();
149+
// } catch {
150+
// continue;
151+
// }
152+
153+
// dbtInstalledPythonPath.push(candidatePythonPath);
154+
// }
155+
156+
return (this.executionDetails = {
157+
getPythonPath: () => {
158+
if (dbtInstalledPythonPath.length > 0) {
159+
return dbtInstalledPythonPath[0];
160+
} else {
161+
return api.settings.getExecutionDetails(workspace.workspaceFile)
162+
.execCommand[0];
163+
}
164+
},
165+
onDidChangeExecutionDetails: api.settings.onDidChangeExecutionDetails,
166+
// There are 3 places from where we can get environments variables:
167+
// 1. process env 2. integrated terminal env 3. dot env file(we can get this from python extension)
168+
// Collecting env vars from all 3 places and merging them into one in the above order
169+
// While merging, also tagging the places from where the env var has come.
170+
getEnvVars: () => {
171+
const envVars: EnvironmentVariables = {};
172+
for (const key in process.env) {
173+
envVars[key] = process.env[key];
174+
this.environmentVariableSource[key] = "process";
175+
}
176+
try {
177+
const integratedEnv:
178+
| Record<string, Record<string, string>>
179+
| undefined = workspace
180+
.getConfiguration("terminal")
181+
.get("integrated.env");
182+
if (integratedEnv) {
183+
// parse vs code environment variables
184+
for (const prop in integratedEnv) {
185+
// Ignore any settings not supported by the terminal
186+
// We don't know which os is used in terminal unfortunately, so we just merge all of them.
187+
if (!["osx", "windows", "linux"].includes(prop)) {
188+
this.dbtTerminal.debug(
189+
"pythonEnvironment:envVars",
190+
"Loading env vars from config.terminal.integrated.env",
191+
"Ignoring invalid property " + prop,
192+
);
193+
continue;
194+
}
195+
this.dbtTerminal.debug(
196+
"pythonEnvironment:envVars",
197+
"Loading env vars from config.terminal.integrated.env",
198+
"Merging from " + prop,
199+
Object.keys(integratedEnv[prop]),
200+
);
201+
for (const key in integratedEnv[prop]) {
202+
envVars[key] = this.substituteSettingsVariables(
203+
integratedEnv[prop][key],
204+
process.env,
205+
);
206+
this.environmentVariableSource[key] = "integrated";
207+
}
208+
}
209+
}
210+
if (api.environment) {
211+
const workspacePath = workspace.workspaceFolders![0];
212+
this.dbtTerminal.debug(
213+
"pythonEnvironment:envVars",
214+
`workspacePath:${workspacePath.uri.fsPath}`,
215+
);
216+
const workspaceEnv =
217+
api.environments.getEnvironmentVariables(workspacePath);
218+
this.dbtTerminal.debug(
219+
"pythonEnvironment:envVars",
220+
`workspaceEnv:${Object.keys(workspaceEnv)}`,
221+
);
222+
for (const key in workspaceEnv) {
223+
// env var from python extension also includes env var from process env
224+
// therefore only merging those env var, which are not present in process env
225+
// or whose value differ from process env
226+
if (!(key in envVars) || workspaceEnv[key] !== envVars[key]) {
227+
envVars[key] = workspaceEnv[key];
228+
this.environmentVariableSource[key] = "dotenv";
229+
}
230+
}
231+
}
232+
} catch (e: any) {
233+
this.dbtTerminal.error(
234+
"getEnvVarsError",
235+
"Could not call environment api",
236+
e,
237+
);
238+
}
239+
240+
return envVars;
241+
},
242+
});
243+
}
244+
}

‎src/manifest/runtimePythonEnvironmentProvider.ts

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
import { injectable, inject } from "inversify";
2+
import { PythonEnvironment } from "./pythonEnvironment";
3+
import {
4+
PythonEnvironmentProvider,
5+
RuntimePythonEnvironment,
6+
} from "../dbt_integration/pythonEnvironment";
7+
8+
@injectable()
9+
export class VSCodeRuntimePythonEnvironmentProvider
10+
implements PythonEnvironmentProvider
11+
{
12+
private callbacks: ((environment: RuntimePythonEnvironment) => void)[] = [];
13+
14+
constructor(
15+
@inject(PythonEnvironment)
16+
private vscodeEnvironment: PythonEnvironment,
17+
) {
18+
// TODO: Re-implement python environment change handling
19+
// Need to access the VSCode extension's python environment changes
20+
// This would require the VSCode layer to notify this provider
21+
}
22+
23+
getCurrentEnvironment(): RuntimePythonEnvironment {
24+
return {
25+
pythonPath: this.vscodeEnvironment.pythonPath,
26+
environmentVariables: this.vscodeEnvironment.environmentVariables,
27+
};
28+
}
29+
30+
onEnvironmentChanged(
31+
callback: (environment: RuntimePythonEnvironment) => void,
32+
): void {
33+
this.callbacks.push(callback);
34+
}
35+
}
36+
37+
@injectable()
38+
export class StaticRuntimePythonEnvironment
39+
implements RuntimePythonEnvironment
40+
{
41+
constructor(
42+
@inject(PythonEnvironment)
43+
private vscodeEnvironment: PythonEnvironment,
44+
) {}
45+
46+
get pythonPath(): string {
47+
return this.vscodeEnvironment.pythonPath;
48+
}
49+
50+
get environmentVariables() {
51+
return this.vscodeEnvironment.environmentVariables;
52+
}
53+
}

0 commit comments

Comments
 (0)