Add PYTHONPATH to environment variables if the Python extension makes…

… it available

.vscode/settings.json

@@ -38,6 +38,7 @@
     "MNIST",
     "nfiles",
     "pseudoterminal",
+    "PYTHONPATH",
     "rwlock",
     "sonarjs",
     "spyable",

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,
     {

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,
     {

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,
       {

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()
+    })
   }
 }

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,

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
   }
 }

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,
     {
@@ -314,5 +317,21 @@ describe('CliReader', () => {
 
       await expect(dvcReader.version(cwd)).rejects.toBeTruthy()
     })
+
+    it('should add PYTHONPATH to the environment if it is provided', 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'
+      })
+    })
   })
 })

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)
   }

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))

extension/src/cli/dvc/viewer.ts

@@ -88,12 +88,13 @@ export class DvcViewer extends Disposable implements ICli {
   }
 
   private getOptions(cwd: string, args: Args) {
-    return getOptions(
-      this.config.getPythonBinPath(),
-      this.config.getCliPath(),
+    return getOptions({
+      PYTHONPATH: this.config.getPYTHONPATH(),
+      args: [...args],
+      cliPath: this.config.getCliPath(),
       cwd,
-      ...args
-    )
+      pythonBinPath: this.config.getPythonBinPath()
+    })
   }
 
   private getRunningProcess(cwd: string, ...args: Args) {

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(() => {

extension/src/extensions/python.ts

@@ -11,27 +11,36 @@ interface Settings {
   }
 }
 
+type EnvironmentVariables = { readonly [key: string]: string | undefined }
+type EnvironmentVariablesChangeEvent = {
+  readonly env: EnvironmentVariables
+}
+
 export interface VscodePython {
   ready: Thenable<void>
   settings: Settings
+  environments: {
+    onDidEnvironmentVariablesChange: Event<EnvironmentVariablesChangeEvent>
+    getEnvironmentVariables(): EnvironmentVariables
+  }
 }
 
-const getPythonExtensionSettings = async (): Promise<Settings | undefined> => {
+const getPythonExtensionAPI = async (): Promise<VscodePython | undefined> => {
   const api = await getExtensionAPI<VscodePython>(PYTHON_EXTENSION_ID)
   if (!api) {
     return
   }
   try {
     await api.ready
   } catch {}
-  return api.settings
+  return api
 }
 
 export const getPythonExecutionDetails = async (): Promise<
   string[] | undefined
 > => {
-  const settings = await getPythonExtensionSettings()
-  return settings?.getExecutionDetails().execCommand
+  const api = await getPythonExtensionAPI()
+  return api?.settings?.getExecutionDetails().execCommand
 }
 
 export const getPythonBinPath = async (): Promise<string | undefined> => {
@@ -42,9 +51,19 @@ export const getPythonBinPath = async (): Promise<string | undefined> => {
   }
 }
 
+export const getPYTHONPATH = async (): Promise<string | undefined> => {
+  const api = await getPythonExtensionAPI()
+  return api?.environments?.getEnvironmentVariables().PYTHONPATH
+}
+
 export const getOnDidChangePythonExecutionDetails = async () => {
-  const settings = await getPythonExtensionSettings()
-  return settings?.onDidChangeExecutionDetails
+  const api = await getPythonExtensionAPI()
+  return api?.settings?.onDidChangeExecutionDetails
+}
+
+export const getOnDidChangePythonEnvironmentVariables = async () => {
+  const api = await getPythonExtensionAPI()
+  return api?.environments?.onDidEnvironmentVariablesChange
 }
 
 export const isPythonExtensionInstalled = () => isInstalled(PYTHON_EXTENSION_ID)