Add Python Extension's "Create Env" command into DVC Setup

extension/src/extensions/python.ts

@@ -68,6 +68,10 @@ export const getOnDidChangePythonEnvironmentVariables = async () => {
 
 export const isPythonExtensionInstalled = () => isInstalled(PYTHON_EXTENSION_ID)
 
+export const createPythonEnv = () => {
+  void commands.executeCommand('python.createEnvironment')
+}
+
 export const selectPythonInterpreter = () => {
   void commands.executeCommand('python.setInterpreter')
 }

extension/src/setup/index.ts

@@ -13,7 +13,11 @@ import { validateTokenInput } from './inputBox'
 import { findPythonBinForInstall } from './autoInstall'
 import { run, runWithRecheck, runWorkspace } from './runner'
 import { isStudioAccessToken } from './token'
-import { pickFocusedProjects } from './quickPick'
+import {
+  PYTHON_EXTENSION_ACTION,
+  pickFocusedProjects,
+  pickPythonExtensionAction
+} from './quickPick'
 import { ViewKey } from '../webview/constants'
 import { BaseRepository } from '../webview/repository'
 import { Resource } from '../resourceLocator'
@@ -56,6 +60,7 @@ import { Title } from '../vscode/title'
 import { getDVCAppDir } from '../util/appdirs'
 import { getOptions } from '../cli/dvc/options'
 import { isAboveLatestTestedVersion } from '../cli/dvc/version'
+import { createPythonEnv, selectPythonInterpreter } from '../extensions/python'
 
 export class Setup
   extends BaseRepository<TSetupData>
@@ -412,7 +417,8 @@ export class Setup
     const webviewMessages = new WebviewMessages(
       () => this.getWebview(),
       () => this.initializeGit(),
-      (offline: boolean) => this.updateStudioOffline(offline)
+      (offline: boolean) => this.updateStudioOffline(offline),
+      () => this.updatePythonEnvironment()
     )
     this.dispose.track(
       this.onDidReceivedWebviewMessage(message =>
@@ -504,6 +510,18 @@ export class Setup
     }
   }
 
+  private async updatePythonEnvironment() {
+    const value = await pickPythonExtensionAction()
+
+    if (!value) {
+      return
+    }
+
+    return value === PYTHON_EXTENSION_ACTION.CREATE_ENV
+      ? createPythonEnv()
+      : selectPythonInterpreter()
+  }
+
   private needsGitCommit(needsGitInit: boolean) {
     if (needsGitInit) {
       return true

extension/src/setup/quickPick.test.ts

@@ -1,12 +1,17 @@
 import { join } from 'path'
-import { pickFocusedProjects } from './quickPick'
-import { quickPickManyValues } from '../vscode/quickPick'
+import {
+  PYTHON_EXTENSION_ACTION,
+  pickFocusedProjects,
+  pickPythonExtensionAction
+} from './quickPick'
+import { quickPickManyValues, quickPickValue } from '../vscode/quickPick'
 import { Toast } from '../vscode/toast'
 import { Title } from '../vscode/title'
 
 jest.mock('../vscode/quickPick')
 
 const mockedQuickPickManyValues = jest.mocked(quickPickManyValues)
+const mockedQuickValue = jest.mocked(quickPickValue)
 
 const mockedToast = jest.mocked(Toast)
 const mockedShowError = jest.fn()
@@ -55,3 +60,28 @@ describe('pickFocusedProjects', () => {
     )
   })
 })
+
+describe('pickPythonExtensionAction', () => {
+  it('should call a quick pick with the correct values', () => {
+    void pickPythonExtensionAction()
+
+    expect(mockedQuickValue).toHaveBeenCalledWith(
+      [
+        {
+          description: 'Create an environment',
+          label: 'Create',
+          value: PYTHON_EXTENSION_ACTION.CREATE_ENV
+        },
+        {
+          description: 'Choose from already created environments',
+          label: 'Select',
+          value: PYTHON_EXTENSION_ACTION.SET_INTERPRETER
+        }
+      ],
+      {
+        placeHolder: 'Select or Create a Python Environment',
+        title: Title.UPDATE_PYTHON_ENVIRONMENT
+      }
+    )
+  })
+})

extension/src/setup/quickPick.ts

@@ -1,8 +1,13 @@
 import { sortCollectedArray } from '../util/array'
-import { quickPickManyValues } from '../vscode/quickPick'
+import { quickPickManyValues, quickPickValue } from '../vscode/quickPick'
 import { Title } from '../vscode/title'
 import { Toast } from '../vscode/toast'
 
+export enum PYTHON_EXTENSION_ACTION {
+  CREATE_ENV = 1,
+  SET_INTERPRETER = 2
+}
+
 export const pickFocusedProjects = async (
   projects: string[],
   currentProjects: string[]
@@ -26,3 +31,24 @@ export const pickFocusedProjects = async (
 
   return sortCollectedArray(values)
 }
+
+export const pickPythonExtensionAction = (): Thenable<
+  PYTHON_EXTENSION_ACTION | undefined
+> => {
+  const options = [
+    {
+      description: 'Create an environment',
+      label: 'Create',
+      value: PYTHON_EXTENSION_ACTION.CREATE_ENV
+    },
+    {
+      description: 'Choose from already created environments',
+      label: 'Select',
+      value: PYTHON_EXTENSION_ACTION.SET_INTERPRETER
+    }
+  ]
+  return quickPickValue<PYTHON_EXTENSION_ACTION>(options, {
+    placeHolder: 'Select or Create a Python Environment',
+    title: Title.UPDATE_PYTHON_ENVIRONMENT
+  })
+}

extension/src/setup/webview/messages.ts

@@ -8,7 +8,6 @@ import {
 import { BaseWebview } from '../../webview'
 import { sendTelemetryEvent } from '../../telemetry'
 import { EventName } from '../../telemetry/constants'
-import { selectPythonInterpreter } from '../../extensions/python'
 import { autoInstallDvc, autoUpgradeDvc } from '../autoInstall'
 import {
   RegisteredCliCommands,
@@ -20,15 +19,18 @@ export class WebviewMessages {
   private readonly getWebview: () => BaseWebview<TSetupData> | undefined
   private readonly initializeGit: () => void
   private readonly updateStudioOffline: (offline: boolean) => Promise<void>
+  private readonly updatePythonEnv: () => Promise<void>
 
   constructor(
     getWebview: () => BaseWebview<TSetupData> | undefined,
     initializeGit: () => void,
-    updateStudioOffline: (shareLive: boolean) => Promise<void>
+    updateStudioOffline: (shareLive: boolean) => Promise<void>,
+    updatePythonEnv: () => Promise<void>
   ) {
     this.getWebview = getWebview
     this.initializeGit = initializeGit
     this.updateStudioOffline = updateStudioOffline
+    this.updatePythonEnv = updatePythonEnv
   }
 
   public sendWebviewMessage({
@@ -79,8 +81,8 @@ export class WebviewMessages {
         return commands.executeCommand(RegisteredCommands.EXTENSION_GET_STARTED)
       case MessageFromWebviewType.SHOW_SCM_PANEL:
         return this.showScmForCommit()
-      case MessageFromWebviewType.SELECT_PYTHON_INTERPRETER:
-        return this.selectPythonInterpreter()
+      case MessageFromWebviewType.UPDATE_PYTHON_ENVIRONMENT:
+        return this.updatePythonEnvironment()
       case MessageFromWebviewType.INSTALL_DVC:
         return this.installDvc()
       case MessageFromWebviewType.UPGRADE_DVC:
@@ -131,13 +133,13 @@ export class WebviewMessages {
     return commands.executeCommand('workbench.view.scm')
   }
 
-  private selectPythonInterpreter() {
+  private updatePythonEnvironment() {
     sendTelemetryEvent(
-      EventName.VIEWS_SETUP_SELECT_PYTHON_INTERPRETER,
+      EventName.VIEWS_SETUP_UPDATE_PYTHON_ENVIRONMENT,
       undefined,
       undefined
     )
-    return selectPythonInterpreter()
+    return this.updatePythonEnv()
   }
 
   private upgradeDvc() {

extension/src/telemetry/constants.ts

@@ -85,9 +85,9 @@ export const EventName = Object.assign(
     VIEWS_SETUP_FOCUS_CHANGED: 'views.setup.focusChanged',
     VIEWS_SETUP_INIT_GIT: 'views.setup.initializeGit',
     VIEWS_SETUP_INSTALL_DVC: 'views.setup.installDvc',
-    VIEWS_SETUP_SELECT_PYTHON_INTERPRETER:
-      'views.setup.selectPythonInterpreter',
     VIEWS_SETUP_SHOW_SCM_FOR_COMMIT: 'views.setup.showScmForCommit',
+    VIEWS_SETUP_UPDATE_PYTHON_ENVIRONMENT:
+      'views.setup.updatePythonEnvironment',
     VIEWS_SETUP_UPGRADE_DVC: 'view.setup.upgradeDvc',
 
     VIEWS_TERMINAL_CLOSED: 'views.terminal.closed',
@@ -277,7 +277,7 @@ export interface IEventNamePropertyMapping {
   [EventName.VIEWS_SETUP_CLOSE]: undefined
   [EventName.VIEWS_SETUP_CREATED]: undefined
   [EventName.VIEWS_SETUP_FOCUS_CHANGED]: undefined
-  [EventName.VIEWS_SETUP_SELECT_PYTHON_INTERPRETER]: undefined
+  [EventName.VIEWS_SETUP_UPDATE_PYTHON_ENVIRONMENT]: undefined
   [EventName.VIEWS_SETUP_SHOW_SCM_FOR_COMMIT]: undefined
   [EventName.VIEWS_SETUP_INIT_GIT]: undefined
   [EventName.VIEWS_SETUP_INSTALL_DVC]: undefined

extension/src/test/suite/setup/index.test.ts

@@ -48,6 +48,7 @@ import { SetupSection } from '../../../setup/webview/contract'
 import { getFirstWorkspaceFolder } from '../../../vscode/workspaceFolders'
 import { Response } from '../../../vscode/response'
 import { DvcConfig } from '../../../cli/dvc/config'
+import * as QuickPickUtil from '../../../setup/quickPick'
 
 suite('Setup Test Suite', () => {
   const disposable = Disposable.fn()
@@ -156,21 +157,79 @@ suite('Setup Test Suite', () => {
       expect(mockAutoUpgradeDvc).to.be.calledOnce
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
-    it('should handle a select Python interpreter message from the webview', async () => {
-      const { messageSpy, mockExecuteCommand, setup } = buildSetup(disposable)
-      const setInterpreterCommand = 'python.setInterpreter'
+    it('should handle an update Python environment message from the webview', async () => {
+      const { messageSpy, setup, mockExecuteCommand } = buildSetup(disposable)
 
       const webview = await setup.showWebview()
       await webview.isReady()
 
+      const mockPickExtensionAction = stub(
+        QuickPickUtil,
+        'pickPythonExtensionAction'
+      )
+
+      const firstQuickPickEvent = new Promise(resolve => {
+        mockPickExtensionAction.onFirstCall().callsFake(() => {
+          resolve(undefined)
+          return Promise.resolve(undefined)
+        })
+      })
+
       const mockMessageReceived = getMessageReceivedEmitter(webview)
 
       messageSpy.resetHistory()
       mockMessageReceived.fire({
-        type: MessageFromWebviewType.SELECT_PYTHON_INTERPRETER
+        type: MessageFromWebviewType.UPDATE_PYTHON_ENVIRONMENT
+      })
+
+      await firstQuickPickEvent
+
+      expect(mockExecuteCommand).to.not.be.calledWithExactly(
+        'python.setInterpreter'
+      )
+      expect(mockExecuteCommand).to.not.be.calledWithExactly(
+        'python.createEnvironment'
+      )
+
+      const secondQuickPickEvent = new Promise(resolve => {
+        mockPickExtensionAction.onSecondCall().callsFake(() => {
+          resolve(undefined)
+          return Promise.resolve(
+            QuickPickUtil.PYTHON_EXTENSION_ACTION.CREATE_ENV
+          )
+        })
+      })
+
+      messageSpy.resetHistory()
+      mockMessageReceived.fire({
+        type: MessageFromWebviewType.UPDATE_PYTHON_ENVIRONMENT
+      })
+
+      await secondQuickPickEvent
+
+      expect(mockExecuteCommand).to.be.calledWithExactly(
+        'python.createEnvironment'
+      )
+
+      const thirdQuickPickEvent = new Promise(resolve => {
+        mockPickExtensionAction.onThirdCall().callsFake(() => {
+          resolve(undefined)
+          return Promise.resolve(
+            QuickPickUtil.PYTHON_EXTENSION_ACTION.SET_INTERPRETER
+          )
+        })
+      })
+
+      messageSpy.resetHistory()
+      mockMessageReceived.fire({
+        type: MessageFromWebviewType.UPDATE_PYTHON_ENVIRONMENT
       })
 
-      expect(mockExecuteCommand).to.be.calledWithExactly(setInterpreterCommand)
+      await thirdQuickPickEvent
+
+      expect(mockExecuteCommand).to.be.calledWithExactly(
+        'python.setInterpreter'
+      )
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should handle a show source control panel message from the webview', async () => {

extension/src/vscode/title.ts

@@ -41,7 +41,8 @@ export enum Title {
   SELECT_TRAINING_SCRIPT = 'Select your Training Script',
   SETUP_WORKSPACE = 'Setup the Workspace',
   SET_EXPERIMENTS_HEADER_HEIGHT = 'Set the Maximum Experiment Table Header Height',
-  SET_REMOTE_AS_DEFAULT = 'Set Default Remote'
+  SET_REMOTE_AS_DEFAULT = 'Set Default Remote',
+  UPDATE_PYTHON_ENVIRONMENT = 'Update Python Environment (selected with the Python Extension)'
 }
 
 export const getEnterValueTitle = (path: string): Title =>

extension/src/webview/contract.ts

@@ -47,7 +47,6 @@ export enum MessageFromWebviewType {
   SELECT_EXPERIMENTS = 'select-experiments',
   SELECT_COLUMNS = 'select-columns',
   SELECT_PLOTS = 'select-plots',
-  SELECT_PYTHON_INTERPRETER = 'select-python-interpreter',
   SET_EXPERIMENTS_FOR_PLOTS = 'set-experiments-for-plots',
   SET_EXPERIMENTS_AND_OPEN_PLOTS = 'set-experiments-and-open-plots',
   SET_STUDIO_SHARE_EXPERIMENTS_LIVE = 'set-studio-share-experiments-live',
@@ -66,6 +65,7 @@ export enum MessageFromWebviewType {
   INITIALIZE_GIT = 'initialize-git',
   SHOW_SCM_PANEL = 'show-scm-panel',
   INSTALL_DVC = 'install-dvc',
+  UPDATE_PYTHON_ENVIRONMENT = 'update-python-environment',
   UPGRADE_DVC = 'upgrade-dvc',
   SETUP_WORKSPACE = 'setup-workspace',
   ZOOM_PLOT = 'zoom-plot',
@@ -197,7 +197,7 @@ export type MessageFromWebview =
     }
   | { type: MessageFromWebviewType.INITIALIZED }
   | { type: MessageFromWebviewType.SELECT_EXPERIMENTS }
-  | { type: MessageFromWebviewType.SELECT_PYTHON_INTERPRETER }
+  | { type: MessageFromWebviewType.UPDATE_PYTHON_ENVIRONMENT }
   | { type: MessageFromWebviewType.SELECT_PLOTS }
   | { type: MessageFromWebviewType.REFRESH_REVISIONS }
   | { type: MessageFromWebviewType.SELECT_COLUMNS }