Add PYTHONPATH to environment variables if the Python extension makes it available (#4045)

mattseddon · web-flow · commit 860a59a8a4c2 · 2023-06-08T23:50:43.000Z
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -38,6 +38,7 @@
     "MNIST",
     "nfiles",
     "pseudoterminal",
+    "PYTHONPATH",
     "rwlock",
     "sonarjs",
     "spyable",
diff --git a/README.md b/README.md
@@ -160,13 +160,22 @@ These are the VS Code [settings] available for the Extension:
 > **Note** that the `Setup The Workspace` command helps you set up the basic
 > ones at the [Workspace level] (saved to `.vscode/setting.json`).
 
+### Python
+
+This extension is integrated with Microsoft's [Python extension]. When possible,
+the Python extension's selected interpreter will be used to locate DVC. The
+`PYTHONPATH` environment variable identified via the [python.envFile] config
+setting is also respected.
+
 [python extension]:
   https://marketplace.visualstudio.com/items?itemName=ms-python.python
 [studio.token]:
   https://dvc.org/doc/user-guide/project-structure/configuration#studio
 [Studio]: https://studio.iterative.ai
 [workspace level]:
   https://code.visualstudio.com/docs/getstarted/settings#_workspace-settings
+[python.envFile]:
+  https://code.visualstudio.com/docs/python/environments#_use-of-the-pythonpath-variable
 
 ## Debugging
 
diff --git a/extension/src/cli/dvc/config.test.ts b/extension/src/cli/dvc/config.test.ts
@@ -42,6 +42,7 @@ describe('DvcConfig', () => {
   const dvcConfig = new DvcConfig(
     {
       getCliPath: () => undefined,
+      getPYTHONPATH: () => undefined,
       getPythonBinPath: () => undefined
     } as unknown as Config,
     {
diff --git a/extension/src/cli/dvc/executor.test.ts b/extension/src/cli/dvc/executor.test.ts
@@ -48,6 +48,7 @@ describe('CliExecutor', () => {
   const dvcExecutor = new DvcExecutor(
     {
       getCliPath: () => undefined,
+      getPYTHONPATH: () => undefined,
       getPythonBinPath: () => undefined
     } as unknown as Config,
     {
diff --git a/extension/src/cli/dvc/index.test.ts b/extension/src/cli/dvc/index.test.ts
@@ -66,6 +66,7 @@ describe('executeDvcProcess', () => {
     const cli = new DvcCli(
       {
         getCliPath: () => undefined,
+        getPYTHONPATH: () => undefined,
         getPythonBinPath: () => undefined
       } as unknown as Config,
       {
@@ -111,6 +112,7 @@ describe('executeDvcProcess', () => {
     const cli = new DvcCli(
       {
         getCliPath: () => '/some/path/to/dvc',
+        getPYTHONPATH: () => undefined,
         getPythonBinPath: () => pythonBinPath
       } as unknown as Config,
       {
diff --git a/extension/src/cli/dvc/index.ts b/extension/src/cli/dvc/index.ts
@@ -35,11 +35,12 @@ export class DvcCli extends Cli {
   }
 
   protected getOptions(cwd: string, ...args: Args) {
-    return getOptions(
-      this.extensionConfig.getPythonBinPath(),
-      this.extensionConfig.getCliPath(),
+    return getOptions({
+      PYTHONPATH: this.extensionConfig.getPYTHONPATH(),
+      args: [...args],
+      cliPath: this.extensionConfig.getCliPath(),
       cwd,
-      ...args
-    )
+      pythonBinPath: this.extensionConfig.getPythonBinPath()
+    })
   }
 }
diff --git a/extension/src/cli/dvc/options.test.ts b/extension/src/cli/dvc/options.test.ts
@@ -25,7 +25,13 @@ describe('getOptions', () => {
   const cwd = join('path', 'to', 'work', 'dir')
 
   it('should give the correct options given a basic environment', () => {
-    const options = getOptions(undefined, '', cwd, Command.CHECKOUT, Flag.FORCE)
+    const options = getOptions({
+      PYTHONPATH: undefined,
+      args: [Command.CHECKOUT, Flag.FORCE],
+      cliPath: '',
+      cwd,
+      pythonBinPath: undefined
+    })
     expect(options).toStrictEqual({
       args: ['checkout', '-f'],
       cwd,
@@ -36,7 +42,13 @@ describe('getOptions', () => {
 
   it('should append -m dvc to the args and use the python as the executable if only an isolated python env is in use', () => {
     const pythonBinPath = join('path', 'to', 'python', '.venv', 'python')
-    const options = getOptions(pythonBinPath, '', cwd, Command.PULL)
+    const options = getOptions({
+      PYTHONPATH: undefined,
+      args: [Command.PULL],
+      cliPath: '',
+      cwd,
+      pythonBinPath
+    })
     expect(options).toStrictEqual({
       args: ['-m', 'dvc', 'pull'],
       cwd,
@@ -53,7 +65,13 @@ describe('getOptions', () => {
   it('should append to the PATH but only use the path to the cli if both an isolated python env and path to dvc are in use', () => {
     const pythonBinPath = join('path', 'to', 'python', '.venv', 'python')
     const cliPath = join('custom', 'path', 'to', 'dvc')
-    const options = getOptions(pythonBinPath, cliPath, cwd, Command.PULL)
+    const options = getOptions({
+      PYTHONPATH: undefined,
+      args: [Command.PULL],
+      cliPath,
+      cwd,
+      pythonBinPath
+    })
     expect(options).toStrictEqual({
       args: ['pull'],
       cwd,
@@ -66,4 +84,30 @@ describe('getOptions', () => {
       executable: cliPath
     })
   })
+
+  it('should add PYTHONPATH to the environment variables if it is provided', () => {
+    const PYTHONPATH = join('path', 'to', 'project')
+    const venvPath = join(PYTHONPATH, '.venv')
+    const pythonBinPath = join(venvPath, 'python')
+    const cliPath = ''
+    const options = getOptions({
+      PYTHONPATH,
+      args: [Command.PULL],
+      cliPath,
+      cwd,
+      pythonBinPath
+    })
+    expect(options).toStrictEqual({
+      args: ['-m', 'dvc', 'pull'],
+      cwd,
+      env: {
+        DVCLIVE_OPEN: 'false',
+        DVC_NO_ANALYTICS: 'true',
+        GIT_TERMINAL_PROMPT: '0',
+        PATH: joinEnvPath(venvPath, mockedPATH),
+        PYTHONPATH
+      },
+      executable: pythonBinPath
+    })
+  })
 })
diff --git a/extension/src/cli/dvc/options.ts b/extension/src/cli/dvc/options.ts
@@ -14,16 +14,23 @@ const getPATH = (existingPath: string, pythonBinPath?: string): string => {
   return joinEnvPath(python, existingPath)
 }
 
-const getEnv = (pythonBinPath?: string): NodeJS.ProcessEnv => {
-  const env = getProcessEnv()
-  const PATH = getPATH(env?.PATH as string, pythonBinPath)
-  return {
-    ...env,
+const getEnv = (
+  pythonBinPath: string | undefined,
+  PYTHONPATH: string | undefined
+): NodeJS.ProcessEnv => {
+  const existingEnv = getProcessEnv()
+  const PATH = getPATH(existingEnv?.PATH as string, pythonBinPath)
+  const env: NodeJS.ProcessEnv = {
+    ...existingEnv,
     DVCLIVE_OPEN: 'false',
     DVC_NO_ANALYTICS: 'true',
     GIT_TERMINAL_PROMPT: '0',
     PATH
   }
+  if (PYTHONPATH) {
+    env.PYTHONPATH = PYTHONPATH
+  }
+  return env
 }
 
 const getArgs = (
@@ -35,18 +42,24 @@ const getArgs = (
 const getExecutable = (pythonBinPath: string | undefined, cliPath: string) =>
   cliPath || pythonBinPath || 'dvc'
 
-export const getOptions = (
-  pythonBinPath: string | undefined,
-  cliPath: string,
-  cwd: string,
-  ...originalArgs: Args
-): ExecutionOptions => {
+export const getOptions = ({
+  args = [],
+  cliPath,
+  cwd,
+  pythonBinPath,
+  PYTHONPATH
+}: {
+  args?: Args
+  cliPath: string
+  cwd: string
+  pythonBinPath: string | undefined
+  PYTHONPATH: string | undefined
+}): ExecutionOptions => {
   const executable = getExecutable(pythonBinPath, cliPath)
-  const args = getArgs(pythonBinPath, cliPath, ...originalArgs)
   return {
-    args,
+    args: getArgs(pythonBinPath, cliPath, ...args),
     cwd: getCaseSensitiveCwd(cwd),
-    env: getEnv(pythonBinPath),
+    env: getEnv(pythonBinPath, PYTHONPATH),
     executable
   }
 }
diff --git a/extension/src/cli/dvc/reader.test.ts b/extension/src/cli/dvc/reader.test.ts
@@ -13,6 +13,7 @@ import expShowFixture from '../../test/fixtures/expShow/base/output'
 import plotsDiffFixture from '../../test/fixtures/plotsDiff/output/minimal'
 import { Config } from '../../config'
 import { joinEnvPath } from '../../util/env'
+import { dvcDemoPath } from '../../test/util'
 
 jest.mock('vscode')
 jest.mock('@hediet/std/disposable')
@@ -34,6 +35,7 @@ const mockedEnv = {
 const JSON_FLAG = '--json'
 
 const mockedGetPythonBinPath = jest.fn()
+const mockedGetPYTHONPATH = jest.fn()
 
 beforeEach(() => {
   jest.resetAllMocks()
@@ -53,6 +55,7 @@ describe('CliReader', () => {
   const dvcReader = new DvcReader(
     {
       getCliPath: () => undefined,
+      getPYTHONPATH: mockedGetPYTHONPATH,
       getPythonBinPath: mockedGetPythonBinPath
     } as unknown as Config,
     {
@@ -311,5 +314,21 @@ describe('CliReader', () => {
 
       await expect(dvcReader.version(cwd)).rejects.toBeTruthy()
     })
+
+    it('should add PYTHONPATH to the environment used to create a DVC process when it is available from the extension config', async () => {
+      mockedGetPYTHONPATH.mockReturnValue(dvcDemoPath)
+      const cwd = __dirname
+      const stdout = '3.9.11'
+      mockedCreateProcess.mockReturnValueOnce(getMockedProcess(stdout))
+      const output = await dvcReader.version(cwd)
+
+      expect(output).toStrictEqual(stdout)
+      expect(mockedCreateProcess).toHaveBeenCalledWith({
+        args: ['--version'],
+        cwd,
+        env: { ...mockedEnv, PYTHONPATH: dvcDemoPath },
+        executable: 'dvc'
+      })
+    })
   })
 })
diff --git a/extension/src/cli/dvc/reader.ts b/extension/src/cli/dvc/reader.ts
@@ -124,12 +124,13 @@ export class DvcReader extends DvcCli {
   }
 
   public globalVersion(cwd: string): Promise<string> {
-    const options = getOptions(
-      undefined,
-      this.extensionConfig.getCliPath(),
+    const options = getOptions({
+      PYTHONPATH: undefined,
+      args: [Flag.VERSION],
+      cliPath: this.extensionConfig.getCliPath(),
       cwd,
-      Flag.VERSION
-    )
+      pythonBinPath: undefined
+    })
 
     return this.executeProcess(options)
   }
diff --git a/extension/src/cli/dvc/runner.ts b/extension/src/cli/dvc/runner.ts
@@ -141,12 +141,13 @@ export class DvcRunner extends Disposable implements ICli {
   }
 
   private createProcess({ cwd, args }: { cwd: string; args: Args }): Process {
-    const options = getOptions(
-      this.config.getPythonBinPath(),
-      this.config.getCliPath(),
+    const options = getOptions({
+      PYTHONPATH: this.config.getPYTHONPATH(),
+      args: [...args],
+      cliPath: this.config.getCliPath(),
       cwd,
-      ...args
-    )
+      pythonBinPath: this.config.getPythonBinPath()
+    })
     const command = getCommandString(options)
     const stopWatch = new StopWatch()
     const process = this.dispose.track(createProcess(options))
diff --git a/extension/src/cli/dvc/viewer.ts b/extension/src/cli/dvc/viewer.ts
@@ -77,25 +77,24 @@ export class DvcViewer extends Disposable implements ICli {
   }
 
   private viewProcess(name: string, cwd: string, args: Args) {
+    const options = getOptions({
+      PYTHONPATH: this.config.getPYTHONPATH(),
+      args: [...args],
+      cliPath: this.config.getCliPath(),
+      cwd,
+      pythonBinPath: this.config.getPythonBinPath()
+    })
+
     return this.dispose.track(
       new ViewableCliProcess(
         `DVC: ${name}`,
-        this.getOptions(cwd, args),
+        options,
         this.processStarted,
         this.processCompleted
       )
     )
   }
 
-  private getOptions(cwd: string, args: Args) {
-    return getOptions(
-      this.config.getPythonBinPath(),
-      this.config.getCliPath(),
-      cwd,
-      ...args
-    )
-  }
-
   private getRunningProcess(cwd: string, ...args: Args) {
     return this.processes[this.getId(cwd, ...args)]
   }
diff --git a/extension/src/config.ts b/extension/src/config.ts
@@ -2,8 +2,10 @@ import { EventEmitter, Event, workspace } from 'vscode'
 import { Disposable } from '@hediet/std/disposable'
 import isEqual from 'lodash.isequal'
 import {
+  getOnDidChangePythonEnvironmentVariables,
   getOnDidChangePythonExecutionDetails,
   getPythonBinPath,
+  getPYTHONPATH,
   isPythonExtensionInstalled
 } from './extensions/python'
 import { ConfigKey, getConfigValue, setConfigValue } from './vscode/config'
@@ -20,6 +22,8 @@ export class Config extends DeferredDisposable {
 
   private dvcPath = this.getCliPath()
 
+  private PYTHONPATH: string | undefined
+
   private focusedProjects: string[] | undefined
 
   private readonly configurationDetailsChanged: EventEmitter<void>
@@ -41,6 +45,8 @@ export class Config extends DeferredDisposable {
     void this.onDidChangePythonExecutionDetails()
     this.onDidChangeExtensions()
 
+    void this.watchPYTHONPATH()
+
     this.onDidConfigurationChange()
   }
 
@@ -76,6 +82,10 @@ export class Config extends DeferredDisposable {
     void setConfigValue(ConfigKey.FOCUSED_PROJECTS, focusedProjects)
   }
 
+  public getPYTHONPATH() {
+    return this.PYTHONPATH
+  }
+
   public isPythonExtensionUsed() {
     return !getConfigValue(ConfigKey.PYTHON_PATH) && !!this.pythonBinPath
   }
@@ -99,6 +109,17 @@ export class Config extends DeferredDisposable {
     )
   }
 
+  private async watchPYTHONPATH() {
+    this.PYTHONPATH = await getPYTHONPATH()
+    const onDidChangePythonEnvironmentVariables =
+      await getOnDidChangePythonEnvironmentVariables()
+    this.pythonExecutionDetailsListener = this.dispose.track(
+      onDidChangePythonEnvironmentVariables?.(({ env }) => {
+        this.PYTHONPATH = env.PYTHONPATH
+      })
+    )
+  }
+
   private onDidChangeExtensions() {
     this.dispose.track(
       this.onDidChangeExtensionsEvent(() => {
diff --git a/extension/src/extensions/python.ts b/extension/src/extensions/python.ts
diff --git a/extension/src/setup/index.ts b/extension/src/setup/index.ts
diff --git a/extension/src/test/suite/cli/dvc/runner.test.ts b/extension/src/test/suite/cli/dvc/runner.test.ts
diff --git a/extension/src/test/suite/cli/dvc/viewer.test.ts b/extension/src/test/suite/cli/dvc/viewer.test.ts
diff --git a/extension/src/test/suite/config.test.ts b/extension/src/test/suite/config.test.ts
diff --git a/extension/src/test/suite/status.test.ts b/extension/src/test/suite/status.test.ts

Original file line number	Diff line number	Diff line change
`@@ -42,6 +42,7 @@ describe('DvcConfig', () => {`
`42`	`42`	`const dvcConfig = new DvcConfig(`
`43`	`43`	`{`
`44`	`44`	`getCliPath: () => undefined,`
	`45`	`+ getPYTHONPATH: () => undefined,`
`45`	`46`	`getPythonBinPath: () => undefined`
`46`	`47`	`} as unknown as Config,`
`47`	`48`	`{`
Original file line number	Diff line number	Diff line change
`@@ -48,6 +48,7 @@ describe('CliExecutor', () => {`
`48`	`48`	`const dvcExecutor = new DvcExecutor(`
`49`	`49`	`{`
`50`	`50`	`getCliPath: () => undefined,`
	`51`	`+ getPYTHONPATH: () => undefined,`
`51`	`52`	`getPythonBinPath: () => undefined`
`52`	`53`	`} as unknown as Config,`
`53`	`54`	`{`
Original file line number	Diff line number	Diff line change
`@@ -66,6 +66,7 @@ describe('executeDvcProcess', () => {`
`66`	`66`	`const cli = new DvcCli(`
`67`	`67`	`{`
`68`	`68`	`getCliPath: () => undefined,`
	`69`	`+ getPYTHONPATH: () => undefined,`
`69`	`70`	`getPythonBinPath: () => undefined`
`70`	`71`	`} as unknown as Config,`
`71`	`72`	`{`
`@@ -111,6 +112,7 @@ describe('executeDvcProcess', () => {`
`111`	`112`	`const cli = new DvcCli(`
`112`	`113`	`{`
`113`	`114`	`getCliPath: () => '/some/path/to/dvc',`
	`115`	`+ getPYTHONPATH: () => undefined,`
`114`	`116`	`getPythonBinPath: () => pythonBinPath`
`115`	`117`	`} as unknown as Config,`
`116`	`118`	`{`
Original file line number	Diff line number	Diff line change
`@@ -35,11 +35,12 @@ export class DvcCli extends Cli {`
`35`	`35`	`}`
`36`	`36`
`37`	`37`	`protected getOptions(cwd: string, ...args: Args) {`
`38`		`- return getOptions(`
`39`		`- this.extensionConfig.getPythonBinPath(),`
`40`		`- this.extensionConfig.getCliPath(),`
	`38`	`+ return getOptions({`
	`39`	`+ PYTHONPATH: this.extensionConfig.getPYTHONPATH(),`
	`40`	`+ args: [...args],`
	`41`	`+ cliPath: this.extensionConfig.getCliPath(),`
`41`	`42`	`cwd,`
`42`		`- ...args`
`43`		`- )`
	`43`	`+ pythonBinPath: this.extensionConfig.getPythonBinPath()`
	`44`	`+ })`
`44`	`45`	`}`
`45`	`46`	`}`