Show custom TensorBoard package installation prompt (#14929)

Author: joyceerhl

package.nls.json

@@ -232,6 +232,7 @@
     "TensorBoard.selectAFolderDetail": "Select a log directory containing tfevent files",
     "TensorBoard.selectAnotherFolderDetail": "Use the file explorer to select another folder",
     "TensorBoard.useCurrentWorkingDirectoryDetail": "TensorBoard will search for tfevent files in all subdirectories of the current working directory",
+    "TensorBoard.installPrompt": "The package TensorBoard is required to launch a TensorBoard session. Would you like to install it?",
     "TensorBoard.launchNativeTensorBoardSessionCodeAction": "Launch TensorBoard session",
     "TensorBoard.launchNativeTensorBoardSessionCodeLens": "▶ Launch TensorBoard Session"
 }

src/client/common/installer/productInstaller.ts

@@ -28,7 +28,7 @@ import {
     Product,
     ProductType
 } from '../types';
-import { Common, Installer, Linters } from '../utils/localize';
+import { Common, Installer, Linters, TensorBoard } from '../utils/localize';
 import { isResource, noop } from '../utils/misc';
 import { ProductNames } from './productNames';
 import {
@@ -570,6 +570,20 @@ export class DataScienceInstaller extends BaseInstaller {
     }
 }
 
+export class TensorBoardInstaller extends DataScienceInstaller {
+    protected async promptToInstallImplementation(
+        product: Product,
+        resource: Uri,
+        cancel: CancellationToken
+    ): Promise<InstallerResponse> {
+        // Show a prompt message specific to TensorBoard
+        const yes = Common.bannerLabelYes();
+        const no = Common.bannerLabelNo();
+        const selection = await this.appShell.showErrorMessage(TensorBoard.installPrompt(), ...[yes, no]);
+        return selection === yes ? this.install(product, resource, cancel) : InstallerResponse.Ignore;
+    }
+}
+
 @injectable()
 export class ProductInstaller implements IInstaller {
     private readonly productService: IProductService;
@@ -626,6 +640,8 @@ export class ProductInstaller implements IInstaller {
                 return new RefactoringLibraryInstaller(this.serviceContainer, this.outputChannel);
             case ProductType.DataScience:
                 return new DataScienceInstaller(this.serviceContainer, this.outputChannel);
+            case ProductType.TensorBoard:
+                return new TensorBoardInstaller(this.serviceContainer, this.outputChannel);
             default:
                 break;
         }

src/client/common/installer/productPath.ts

@@ -111,3 +111,13 @@ export class DataScienceProductPathService extends BaseProductPathsService {
         return this.productInstaller.translateProductToModuleName(product, ModuleNamePurpose.run);
     }
 }
+
+@injectable()
+export class TensorBoardProductPathService extends BaseProductPathsService {
+    constructor(@inject(IServiceContainer) serviceContainer: IServiceContainer) {
+        super(serviceContainer);
+    }
+    public getExecutableNameFromSettings(product: Product, _?: Uri): string {
+        return this.productInstaller.translateProductToModuleName(product, ModuleNamePurpose.run);
+    }
+}

src/client/common/installer/productService.ts

@@ -34,7 +34,7 @@ export class ProductService implements IProductService {
         this.ProductTypes.set(Product.nbconvert, ProductType.DataScience);
         this.ProductTypes.set(Product.kernelspec, ProductType.DataScience);
         this.ProductTypes.set(Product.pandas, ProductType.DataScience);
-        this.ProductTypes.set(Product.tensorboard, ProductType.DataScience);
+        this.ProductTypes.set(Product.tensorboard, ProductType.TensorBoard);
     }
     public getProductType(product: Product): ProductType {
         return this.ProductTypes.get(product)!;

src/client/common/installer/serviceRegistry.ts

@@ -18,6 +18,7 @@ import {
     FormatterProductPathService,
     LinterProductPathService,
     RefactoringLibraryProductPathService,
+    TensorBoardProductPathService,
     TestFrameworkProductPathService
 } from './productPath';
 import { ProductService } from './productService';
@@ -75,5 +76,10 @@ export function registerTypes(serviceManager: IServiceManager) {
         DataScienceProductPathService,
         ProductType.DataScience
     );
+    serviceManager.addSingleton<IProductPathService>(
+        IProductPathService,
+        TensorBoardProductPathService,
+        ProductType.TensorBoard
+    );
     serviceManager.addSingleton<IWebviewPanelProvider>(IWebviewPanelProvider, WebviewPanelProvider);
 }

src/client/common/types.ts

@@ -78,7 +78,8 @@ export enum ProductType {
     TestFramework = 'TestFramework',
     RefactoringLibrary = 'RefactoringLibrary',
     WorkspaceSymbols = 'WorkspaceSymbols',
-    DataScience = 'DataScience'
+    DataScience = 'DataScience',
+    TensorBoard = 'TensorBoard'
 }
 
 export enum Product {

src/client/common/utils/localize.ts

@@ -166,6 +166,10 @@ export namespace TensorBoard {
         'TensorBoard.selectAnotherFolderDetail',
         'Use the file explorer to select another folder'
     );
+    export const installPrompt = localize(
+        'TensorBoard.installPrompt',
+        'The package TensorBoard is required to launch a TensorBoard session. Would you like to install it?'
+    );
     export const launchNativeTensorBoardSessionCodeLens = localize(
         'TensorBoard.launchNativeTensorBoardSessionCodeLens',
         '▶ Launch TensorBoard Session'

src/client/tensorBoard/tensorBoardSession.ts

@@ -67,15 +67,15 @@ export class TensorBoardSession {
     // Ensure that the TensorBoard package is installed before we attempt
     // to start a TensorBoard session.
     private async ensureTensorboardIsInstalled() {
-        traceInfo('Ensuring TensorBoard package is installed');
-        if (await this.installer.isInstalled(Product.tensorboard)) {
-            return true;
-        }
+        traceInfo('Ensuring TensorBoard package is installed into active interpreter');
         const interpreter =
             (await this.interpreterService.getActiveInterpreter()) ||
             (await this.commandManager.executeCommand('python.setInterpreter'));
         if (!interpreter) {
-            return undefined;
+            return false;
+        }
+        if (await this.installer.isInstalled(Product.tensorboard, interpreter)) {
+            return true;
         }
         const tokenSource = new CancellationTokenSource();
         const installerToken = tokenSource.token;
@@ -91,7 +91,7 @@ export class TensorBoardSession {
         return response === InstallerResponse.Installed;
     }
 
-    private async showFilePicker() {
+    private async showFilePicker(): Promise<string | undefined> {
         const selection = await window.showOpenDialog({
             canSelectFiles: false,
             canSelectFolders: true,
@@ -102,6 +102,7 @@ export class TensorBoardSession {
         if (selection) {
             return selection[0].fsPath;
         }
+        return undefined;
     }
 
     private getQuickPickItems(logDir: string | undefined) {
@@ -132,23 +133,12 @@ export class TensorBoardSession {
         const selectAFolder = TensorBoard.selectAFolder();
         const selectAnotherFolder = TensorBoard.selectAnotherFolder();
         const items: QuickPickItem[] = this.getQuickPickItems(logDir);
-        const quickPick = window.createQuickPick();
-        quickPick.title = TensorBoard.logDirectoryPrompt();
-        quickPick.canSelectMany = false;
-        quickPick.enabled = false;
-        quickPick.items = items;
-        if (logDir) {
-            quickPick.placeholder = TensorBoard.currentDirectory().format(logDir);
-        }
-        const selection = createDeferred<QuickPickItem>();
-        quickPick.onDidAccept(() => {
-            quickPick.hide();
-            selection.resolve(quickPick.selectedItems[0]);
+        const item = await window.showQuickPick(items, {
+            canPickMany: false,
+            ignoreFocusOut: false,
+            placeHolder: logDir ? TensorBoard.currentDirectory().format(logDir) : undefined
         });
-        quickPick.show();
-        const item = await selection.promise;
-        quickPick.dispose();
-        switch (item.label) {
+        switch (item?.label) {
             case useCurrentWorkingDirectory:
                 return logDir;
             case selectAFolder:
@@ -205,6 +195,7 @@ export class TensorBoardSession {
                 throw new Error(`Timed out after ${timeout / 1000} seconds waiting for TensorBoard to launch.`);
             case 'canceled':
                 traceInfo('Canceled starting TensorBoard session.');
+                observable.dispose();
                 return false;
             case 'success':
                 this.process = observable.proc;

src/test/common/installer.test.ts

@@ -52,6 +52,7 @@ import {
     FormatterProductPathService,
     LinterProductPathService,
     RefactoringLibraryProductPathService,
+    TensorBoardProductPathService,
     TestFrameworkProductPathService
 } from '../../client/common/installer/productPath';
 import { ProductService } from '../../client/common/installer/productService';
@@ -208,6 +209,11 @@ suite('Installer', () => {
             RefactoringLibraryProductPathService,
             ProductType.RefactoringLibrary
         );
+        ioc.serviceManager.addSingleton<IProductPathService>(
+            IProductPathService,
+            TensorBoardProductPathService,
+            ProductType.TensorBoard
+        );
 
         ioc.serviceManager.addSingleton<IActiveResourceService>(IActiveResourceService, ActiveResourceService);
         ioc.serviceManager.addSingleton<IInterpreterPathService>(IInterpreterPathService, InterpreterPathService);
@@ -362,7 +368,8 @@ suite('Installer', () => {
     }
     getNamesAndValues<Product>(Product).forEach((prod) => {
         test(`Ensure install for Product: '${prod.name}' executes the right command in IModuleInstaller`, async function () {
-            if (new ProductService().getProductType(prod.value) === ProductType.DataScience) {
+            const productType = new ProductService().getProductType(prod.value);
+            if (productType === ProductType.DataScience || productType === ProductType.TensorBoard) {
                 // tslint:disable-next-line: no-invalid-this
                 return this.skip();
             }