Open comparison table images when clicking on them (#3467)

* Open comparison table images when clicking on them

* Fix and add tests

* Use join as join function for resourceUrl

* Apply review comments

* Fix tests

* Add getWebviewUri inside of tests

Author: sroy3

extension/src/fileSystem/index.ts

@@ -10,7 +10,7 @@ import {
   writeFileSync
 } from 'fs-extra'
 import { load } from 'js-yaml'
-import { Uri, workspace, window } from 'vscode'
+import { Uri, workspace, window, commands, ViewColumn } from 'vscode'
 import { standardizePath } from './path'
 import { definedAndNonEmpty } from '../util/array'
 import { Logger } from '../common/logger'
@@ -141,6 +141,11 @@ export const openFileInEditor = async (filePath: string) => {
   return document
 }
 
+export const openImageFileInEditor = async (imagePath: string) =>
+  await commands.executeCommand('vscode.open', Uri.file(imagePath), {
+    viewColumn: ViewColumn.Beside
+  })
+
 export const findOrCreateDvcYamlFile = (
   cwd: string,
   trainingScript: string,

extension/src/plots/webview/messages.ts

@@ -21,7 +21,7 @@ import {
 import { PlotsModel } from '../model'
 import { PathsModel } from '../paths/model'
 import { BaseWebview } from '../../webview'
-import { getModifiedTime } from '../../fileSystem'
+import { getModifiedTime, openImageFileInEditor } from '../../fileSystem'
 import { pickCustomPlots, pickMetricAndParam } from '../model/quickPick'
 import { Title } from '../../vscode/title'
 import { ColumnType } from '../../experiments/webview/contract'
@@ -112,9 +112,13 @@ export class WebviewMessages {
       case MessageFromWebviewType.TOGGLE_EXPERIMENT:
         return this.setExperimentStatus(message.payload)
       case MessageFromWebviewType.ZOOM_PLOT:
+        if (message.payload) {
+          const imagePath = this.revertCorrectUrl(message.payload)
+          void openImageFileInEditor(imagePath)
+        }
         return sendTelemetryEvent(
           EventName.VIEWS_PLOTS_ZOOM_PLOT,
-          undefined,
+          { isImage: !!message.payload },
           undefined
         )
       default:
@@ -418,6 +422,15 @@ export class WebviewMessages {
     }
   }
 
+  private revertCorrectUrl(url: string) {
+    const webview = this.getWebview()
+    if (webview) {
+      const toRemove = webview.getWebviewUri('')
+      return url.replace(toRemove, '').split('?')[0]
+    }
+    return url
+  }
+
   private getCheckpointPlots() {
     return this.plots.getCheckpointPlots() || null
   }

extension/src/telemetry/constants.ts

@@ -258,7 +258,7 @@ export interface IEventNamePropertyMapping {
   [EventName.VIEWS_PLOTS_SELECT_EXPERIMENTS]: undefined
   [EventName.VIEWS_PLOTS_SELECT_PLOTS]: undefined
   [EventName.VIEWS_PLOTS_EXPERIMENT_TOGGLE]: undefined
-  [EventName.VIEWS_PLOTS_ZOOM_PLOT]: undefined
+  [EventName.VIEWS_PLOTS_ZOOM_PLOT]: { isImage: boolean }
   [EventName.VIEWS_REORDER_PLOTS_METRICS]: undefined
   [EventName.VIEWS_REORDER_PLOTS_CUSTOM]: undefined
   [EventName.VIEWS_REORDER_PLOTS_TEMPLATES]: undefined

extension/src/test/fixtures/plotsDiff/index.ts

@@ -449,11 +449,8 @@ const getImageData = (baseUrl: string, joinFunc = join) => ({
   ]
 })
 
-export const getOutput = (
-  baseUrl: string,
-  joinFunc?: (...args: string[]) => string
-): PlotsOutput => ({
-  ...getImageData(baseUrl, joinFunc),
+export const getOutput = (baseUrl: string): PlotsOutput => ({
+  ...getImageData(baseUrl),
   ...basicVega,
   ...require('./vega').default
 })
@@ -689,14 +686,18 @@ export const getComparisonWebviewMessage = (
   joinFunc?: (...args: string[]) => string
 ): PlotsComparisonData => {
   const plotAcc = [] as ComparisonPlots
+
   for (const [path, plots] of Object.entries(getImageData(baseUrl, joinFunc))) {
     const revisionsAcc: ComparisonRevisionData = {}
     for (const { url, revisions } of plots) {
       const revision = replaceCommitCLIId(revisions?.[0])
       if (!revision) {
         continue
       }
-      revisionsAcc[revision] = { url: `${url}?${MOCK_IMAGE_MTIME}`, revision }
+      revisionsAcc[revision] = {
+        url: `${url}?${MOCK_IMAGE_MTIME}`,
+        revision
+      }
     }
 
     plotAcc.push({ path, revisions: revisionsAcc })

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

@@ -4,6 +4,7 @@ import cloneDeep from 'lodash.clonedeep'
 import { afterEach, beforeEach, describe, it, suite } from 'mocha'
 import { expect } from 'chai'
 import { restore, spy, stub } from 'sinon'
+import { commands, Uri } from 'vscode'
 import { buildPlots } from '../plots/util'
 import { Disposable } from '../../../extension'
 import expShowFixtureWithoutErrors from '../../fixtures/expShow/base/noErrors'
@@ -529,11 +530,53 @@ suite('Plots Test Suite', () => {
       expect(mockSendTelemetryEvent).to.be.calledOnce
       expect(mockSendTelemetryEvent).to.be.calledWithExactly(
         EventName.VIEWS_PLOTS_ZOOM_PLOT,
-        undefined,
+        { isImage: false },
+        undefined
+      )
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
+
+    it('should handle a plot zoomed message from the webview for an image', async () => {
+      const { plots } = await buildPlots(disposable, plotsDiffFixture)
+
+      const webview = await plots.showWebview()
+
+      const mockSendTelemetryEvent = stub(Telemetry, 'sendTelemetryEvent')
+      const mockMessageReceived = getMessageReceivedEmitter(webview)
+
+      mockMessageReceived.fire({
+        payload: webview.getWebviewUri('a/path.jpg'),
+        type: MessageFromWebviewType.ZOOM_PLOT
+      })
+
+      expect(mockSendTelemetryEvent).to.be.calledOnce
+      expect(mockSendTelemetryEvent).to.be.calledWithExactly(
+        EventName.VIEWS_PLOTS_ZOOM_PLOT,
+        { isImage: true },
         undefined
       )
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
+    it('should open an image when receiving a plot zoomed message from the webview with a payload', async () => {
+      const { plots } = await buildPlots(disposable, plotsDiffFixture)
+
+      const webview = await plots.showWebview()
+      const imagePath = 'some/path/image.jpg'
+
+      stub(Telemetry, 'sendTelemetryEvent')
+      const mockExecuteCommands = stub(commands, 'executeCommand')
+      const mockMessageReceived = getMessageReceivedEmitter(webview)
+
+      mockMessageReceived.fire({
+        payload: webview.getWebviewUri(imagePath),
+        type: MessageFromWebviewType.ZOOM_PLOT
+      })
+
+      expect(mockExecuteCommands).to.be.calledWith(
+        'vscode.open',
+        Uri.file(imagePath)
+      )
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
+
     it('should handle a custom plots reordered message from the webview', async () => {
       const { plots, plotsModel, messageSpy } = await buildPlots(
         disposable,

extension/src/webview/contract.ts

@@ -238,7 +238,7 @@ export type MessageFromWebview =
   | { type: MessageFromWebviewType.OPEN_STUDIO_PROFILE }
   | { type: MessageFromWebviewType.SAVE_STUDIO_TOKEN }
   | { type: MessageFromWebviewType.ADD_CONFIGURATION }
-  | { type: MessageFromWebviewType.ZOOM_PLOT }
+  | { type: MessageFromWebviewType.ZOOM_PLOT; payload?: string }
   | { type: MessageFromWebviewType.OPEN_EXPERIMENTS_WEBVIEW }
 
 export type MessageToWebview<T extends WebviewData> = {

webview/src/plots/components/App.test.tsx

@@ -1436,6 +1436,21 @@ describe('App', () => {
     })
   })
 
+  it('should send a message with the plot path when a comparison table plot is zoomed', () => {
+    renderAppWithOptionalData({
+      comparison: comparisonTableFixture
+    })
+
+    const plot = screen.getAllByTestId('image-plot-button')[0]
+
+    fireEvent.click(plot)
+
+    expect(mockPostMessage).toHaveBeenCalledWith({
+      payload: comparisonTableFixture.plots[0].revisions.workspace.url,
+      type: MessageFromWebviewType.ZOOM_PLOT
+    })
+  })
+
   it('should open a modal with the plot zoomed in when clicking a checkpoint plot', () => {
     renderAppWithOptionalData({
       checkpoint: checkpointPlotsFixture

webview/src/plots/components/comparisonTable/ComparisonTableCell.tsx

@@ -4,6 +4,7 @@ import { ComparisonPlot } from 'dvc/src/plots/webview/contract'
 import styles from './styles.module.scss'
 import { RefreshButton } from '../../../shared/components/button/RefreshButton'
 import { sendMessage } from '../../../shared/vscode'
+import { zoomPlot } from '../messages'
 
 type ComparisonTableCellProps = {
   path: string
@@ -41,10 +42,12 @@ export const ComparisonTableCell: React.FC<ComparisonTableCellProps> = ({
   }
 
   return (
-    <img
-      draggable={false}
-      src={plot.url}
-      alt={`Plot of ${path} (${plot.revision})`}
-    />
+    <button onClick={() => zoomPlot(plot.url)} data-testid="image-plot-button">
+      <img
+        draggable={false}
+        src={plot.url}
+        alt={`Plot of ${path} (${plot.revision})`}
+      />
+    </button>
   )
 }

webview/src/plots/components/messages.ts

@@ -1,5 +1,5 @@
 import { MessageFromWebviewType } from 'dvc/src/webview/contract'
 import { sendMessage } from '../../shared/vscode'
 
-export const zoomPlot = () =>
-  sendMessage({ type: MessageFromWebviewType.ZOOM_PLOT })
+export const zoomPlot = (imagePath?: string) =>
+  sendMessage({ payload: imagePath, type: MessageFromWebviewType.ZOOM_PLOT })