Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 10 commits into from
Jun 14, 2023
35 changes: 34 additions & 1 deletion extension/src/extensions/python.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,8 @@ import { extensions } from 'vscode'
import {
getPythonBinPath,
getOnDidChangePythonExecutionDetails,
VscodePython
VscodePython,
isActivePythonEnvGlobal
} from './python'
import { executeProcess } from '../process/execution'

Expand All @@ -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 = {
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down
21 changes: 21 additions & 0 deletions extension/src/extensions/python.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[I] If we are using this we also need to use onDidChangeActiveEnvironmentPath so that updates are taken into account... unless an update will be triggered in a different way.

[Q] Can we replace our other use of the API with this now?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[I] If we are using this we also need to use onDidChangeActiveEnvironmentPath so that updates are taken into account... unless an update will be triggered in a different way.

Updates are triggered with onDidChangePythonExecutionDetails since that runs when the environment path changes.

[Q] Can we replace our other use of the API with this now?

Apologies, I'm not sure what you mean. What other use of the API could be replaced? Looking at the API, they all seem to be doing different things to me 🤔

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Going to merge this, but happy to update the use of the API in a followup!

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We use getExecutionDetails to get the Python executable for the active environment. Isn't that returned by getActiveEnvironmentPath?

export type EnvironmentPath = {
    /**
     * The ID of the environment.
     */
    readonly id: string;
    /**
     * Path to environment folder or path to python executable that uniquely identifies an environment. Environments
     * lacking a python executable are identified by environment folder paths, whereas other envs can be identified
     * using python executable path.
     */
    readonly path: string;
};

☝🏻 maybe we can't rely on this as it could return a folder and not an executable. Nevermind.

onDidEnvironmentVariablesChange: Event<EnvironmentVariablesChangeEvent>
getEnvironmentVariables(): EnvironmentVariables
}
Expand Down Expand Up @@ -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
Expand Down
114 changes: 83 additions & 31 deletions extension/src/setup/autoInstall.ts
Original file line number Diff line number Diff line change
@@ -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'
Expand All @@ -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)
Expand All @@ -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) {
Expand All @@ -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(
Expand All @@ -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)
}
1 change: 1 addition & 0 deletions extension/src/setup/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -420,6 +420,7 @@ export class Setup
() => this.getWebview(),
() => this.initializeGit(),
(offline: boolean) => this.updateStudioOffline(offline),
() => this.isPythonExtensionUsed(),
() => this.updatePythonEnvironment()
)
this.dispose.track(
Expand Down
17 changes: 12 additions & 5 deletions extension/src/setup/webview/messages.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down Expand Up @@ -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() {
Expand Down
Loading