Add --user flag to global dvc auto installation (#4091)

Author: julieg18

extension/src/extensions/python.test.ts

@@ -2,7 +2,8 @@ import { extensions } from 'vscode'
 import {
   getPythonBinPath,
   getOnDidChangePythonExecutionDetails,
-  VscodePython
+  VscodePython,
+  isActivePythonEnvGlobal
 } from './python'
 import { executeProcess } from '../process/execution'
 
@@ -17,6 +18,7 @@ mockedExtensions.getExtension = mockedGetExtension
 
 const mockedReady = jest.fn()
 const mockedOnDidChangeExecutionDetails = jest.fn()
+const mockedGetActiveEnvironmentPath = jest.fn()
 let mockedExecCommand: string[] | undefined
 
 const mockedSettings = {
@@ -26,7 +28,16 @@ const mockedSettings = {
   onDidChangeExecutionDetails: mockedOnDidChangeExecutionDetails
 }
 
+const mockedEnvironments = {
+  getActiveEnvironmentPath: mockedGetActiveEnvironmentPath,
+  known: [
+    { id: '/usr/bin/python' },
+    { environment: { type: 'VirtualEnvironment' }, id: '/.venv/bin/python' }
+  ]
+}
+
 const mockedVscodePythonAPI = {
+  environments: mockedEnvironments,
   ready: mockedReady,
   settings: mockedSettings
 } as unknown as VscodePython
@@ -63,6 +74,28 @@ describe('getPythonBinPath', () => {
   })
 })
 
+describe('isActivePythonEnvGlobal', () => {
+  it('should return true if active env is global', async () => {
+    mockedGetActiveEnvironmentPath.mockReturnValueOnce({
+      id: '/usr/bin/python'
+    })
+
+    const result = await isActivePythonEnvGlobal()
+
+    expect(result).toStrictEqual(true)
+  })
+
+  it('should return false if active env is not global', async () => {
+    mockedGetActiveEnvironmentPath.mockReturnValueOnce({
+      id: '/.venv/bin/python'
+    })
+
+    const result = await isActivePythonEnvGlobal()
+
+    expect(result).toStrictEqual(false)
+  })
+})
+
 describe('getOnDidChangePythonExecutionDetails', () => {
   it('should return the listener if the python ready promise rejects', async () => {
     mockedReady.mockRejectedValueOnce(undefined)

extension/src/extensions/python.ts

@@ -16,10 +16,19 @@ type EnvironmentVariablesChangeEvent = {
   readonly env: EnvironmentVariables
 }
 
+interface Environment {
+  id: string
+  environment?: {
+    type: string
+  }
+}
+
 export interface VscodePython {
   ready: Thenable<void>
   settings: Settings
   environments: {
+    known: Environment[]
+    getActiveEnvironmentPath: () => { id: string }
     onDidEnvironmentVariablesChange: Event<EnvironmentVariablesChangeEvent>
     getEnvironmentVariables(): EnvironmentVariables
   }
@@ -56,6 +65,18 @@ export const getPYTHONPATH = async (): Promise<string | undefined> => {
   return api?.environments?.getEnvironmentVariables().PYTHONPATH
 }
 
+export const isActivePythonEnvGlobal = async (): Promise<
+  boolean | undefined
+> => {
+  const api = await getPythonExtensionAPI()
+  if (!api?.environments) {
+    return
+  }
+  const envPath = api.environments.getActiveEnvironmentPath()
+  const activeEnv = api.environments.known.find(({ id }) => id === envPath.id)
+  return activeEnv && !activeEnv.environment
+}
+
 export const getOnDidChangePythonExecutionDetails = async () => {
   const api = await getPythonExtensionAPI()
   return api?.settings?.onDidChangeExecutionDetails

extension/src/setup/autoInstall.ts

@@ -1,4 +1,8 @@
-import { getPythonExecutionDetails } from '../extensions/python'
+import { Progress } from 'vscode'
+import {
+  getPythonExecutionDetails,
+  isActivePythonEnvGlobal
+} from '../extensions/python'
 import { findPythonBin, getDefaultPython, installPackages } from '../python'
 import { ConfigKey, getConfigValue } from '../vscode/config'
 import { getFirstWorkspaceFolder } from '../vscode/workspaceFolders'
@@ -16,25 +20,59 @@ export const findPythonBinForInstall = async (): Promise<
   )
 }
 
+const getProcessGlobalArgs = (isGlobal: boolean) => (isGlobal ? ['--user'] : [])
+
+const installPackageAndIncrementProgress = ({
+  root,
+  pythonBinPath,
+  isGlobalEnv,
+  progress,
+  incrementAmount,
+  packageName,
+  successMessage
+}: {
+  root: string
+  pythonBinPath: string
+  isGlobalEnv: boolean
+  progress: Progress<{ message: string; amount: number }>
+  incrementAmount: number
+  packageName: string
+  successMessage: string
+}) =>
+  Toast.runCommandAndIncrementProgress(
+    async () => {
+      await installPackages(
+        root,
+        pythonBinPath,
+        ...getProcessGlobalArgs(isGlobalEnv),
+        packageName
+      )
+      return successMessage
+    },
+    progress,
+    incrementAmount
+  )
+
 const showUpgradeProgress = (
   root: string,
-  pythonBinPath: string
+  pythonBinPath: string,
+  isGlobalEnv: boolean
 ): Thenable<unknown> =>
   Toast.showProgress('Upgrading DVC', async progress => {
     progress.report({ increment: 0 })
 
     progress.report({ increment: 25, message: 'Updating packages...' })
 
     try {
-      await Toast.runCommandAndIncrementProgress(
-        async () => {
-          await installPackages(root, pythonBinPath, 'dvc')
-          return 'Upgraded successfully'
-        },
+      await installPackageAndIncrementProgress({
+        incrementAmount: 75,
+        isGlobalEnv,
+        packageName: 'dvc',
         progress,
-        75
-      )
-
+        pythonBinPath,
+        root,
+        successMessage: 'Upgraded successfully'
+      })
       return Toast.delayProgressClosing()
     } catch (error: unknown) {
       return Toast.reportProgressError(error, progress)
@@ -43,33 +81,36 @@ const showUpgradeProgress = (
 
 const showInstallProgress = (
   root: string,
-  pythonBinPath: string
+  pythonBinPath: string,
+  isGlobalEnv: boolean
 ): Thenable<unknown> =>
   Toast.showProgress('Installing packages', async progress => {
     progress.report({ increment: 0 })
 
     try {
-      await Toast.runCommandAndIncrementProgress(
-        async () => {
-          await installPackages(root, pythonBinPath, 'dvclive')
-          return 'DVCLive Installed'
-        },
+      await installPackageAndIncrementProgress({
+        incrementAmount: 25,
+        isGlobalEnv,
+        packageName: 'dvclive',
         progress,
-        25
-      )
+        pythonBinPath,
+        root,
+        successMessage: 'DVCLive Installed'
+      })
     } catch (error: unknown) {
       return Toast.reportProgressError(error, progress)
     }
 
     try {
-      await Toast.runCommandAndIncrementProgress(
-        async () => {
-          await installPackages(root, pythonBinPath, 'dvc')
-          return 'DVC Installed'
-        },
+      await installPackageAndIncrementProgress({
+        incrementAmount: 75,
+        isGlobalEnv,
+        packageName: 'dvc',
         progress,
-        75
-      )
+        pythonBinPath,
+        root,
+        successMessage: 'DVC Installed'
+      })
 
       return Toast.delayProgressClosing()
     } catch (error: unknown) {
@@ -78,10 +119,17 @@ const showInstallProgress = (
   })
 
 const getArgsAndRunCommand = async (
-  command: (root: string, pythonBinPath: string) => Thenable<unknown>
+  isPythonExtensionUsed: boolean,
+  command: (
+    root: string,
+    pythonBinPath: string,
+    isGlobalEnv: boolean
+  ) => Thenable<unknown>
 ): Promise<unknown> => {
   const pythonBinPath = await findPythonBinForInstall()
   const root = getFirstWorkspaceFolder()
+  const isPythonEnvGlobal =
+    isPythonExtensionUsed && (await isActivePythonEnvGlobal())
 
   if (!root) {
     return Toast.showError(
@@ -95,13 +143,17 @@ const getArgsAndRunCommand = async (
     )
   }
 
-  return command(root, pythonBinPath)
+  return command(root, pythonBinPath, !!isPythonEnvGlobal)
 }
 
-export const autoInstallDvc = (): Promise<unknown> => {
-  return getArgsAndRunCommand(showInstallProgress)
+export const autoInstallDvc = (
+  isPythonExtensionUsed: boolean
+): Promise<unknown> => {
+  return getArgsAndRunCommand(isPythonExtensionUsed, showInstallProgress)
 }
 
-export const autoUpgradeDvc = (): Promise<unknown> => {
-  return getArgsAndRunCommand(showUpgradeProgress)
+export const autoUpgradeDvc = (
+  isPythonExtensionUsed: boolean
+): Promise<unknown> => {
+  return getArgsAndRunCommand(isPythonExtensionUsed, showUpgradeProgress)
 }

extension/src/setup/index.ts

@@ -420,6 +420,7 @@ export class Setup
       () => this.getWebview(),
       () => this.initializeGit(),
       (offline: boolean) => this.updateStudioOffline(offline),
+      () => this.isPythonExtensionUsed(),
       () => this.updatePythonEnvironment()
     )
     this.dispose.track(

extension/src/setup/webview/messages.ts

@@ -8,28 +8,31 @@ import {
 import { BaseWebview } from '../../webview'
 import { sendTelemetryEvent } from '../../telemetry'
 import { EventName } from '../../telemetry/constants'
-import { autoInstallDvc, autoUpgradeDvc } from '../autoInstall'
 import {
   RegisteredCliCommands,
   RegisteredCommands
 } from '../../commands/external'
 import { openUrl } from '../../vscode/external'
+import { autoInstallDvc, autoUpgradeDvc } from '../autoInstall'
 
 export class WebviewMessages {
   private readonly getWebview: () => BaseWebview<TSetupData> | undefined
   private readonly initializeGit: () => void
   private readonly updateStudioOffline: (offline: boolean) => Promise<void>
+  private readonly isPythonExtensionUsed: () => Promise<boolean>
   private readonly updatePythonEnv: () => Promise<void>
 
   constructor(
     getWebview: () => BaseWebview<TSetupData> | undefined,
     initializeGit: () => void,
     updateStudioOffline: (shareLive: boolean) => Promise<void>,
+    isPythonExtensionUsed: () => Promise<boolean>,
     updatePythonEnv: () => Promise<void>
   ) {
     this.getWebview = getWebview
     this.initializeGit = initializeGit
     this.updateStudioOffline = updateStudioOffline
+    this.isPythonExtensionUsed = isPythonExtensionUsed
     this.updatePythonEnv = updatePythonEnv
   }
 
@@ -112,16 +115,20 @@ export class WebviewMessages {
     return this.updatePythonEnv()
   }
 
-  private upgradeDvc() {
+  private async upgradeDvc() {
     sendTelemetryEvent(EventName.VIEWS_SETUP_UPGRADE_DVC, undefined, undefined)
 
-    return autoUpgradeDvc()
+    const isPythonExtensionUsed = await this.isPythonExtensionUsed()
+
+    return autoUpgradeDvc(isPythonExtensionUsed)
   }
 
-  private installDvc() {
+  private async installDvc() {
     sendTelemetryEvent(EventName.VIEWS_SETUP_INSTALL_DVC, undefined, undefined)
 
-    return autoInstallDvc()
+    const isPythonExtensionUsed = await this.isPythonExtensionUsed()
+
+    return autoInstallDvc(isPythonExtensionUsed)
   }
 
   private openStudio() {