Only show plots which are available for selected revisions (#2915)

* only show plots which are available for selected revisions

* use cli id to label mapping in path collection

* filter missing plots from orders (will move to end)

* rework base class

* expose get selected revisions

* test comparison plot ordering with unstable paths

* refactor collection for complexity

* recollect workspace path if mutable revision refetched

* self review

* fix plots welcome screen

Author: mattseddon

extension/src/experiments/columns/model.ts

@@ -4,7 +4,7 @@ import { EXPERIMENT_COLUMN_ID, timestampColumn } from './constants'
 import { Column, ColumnType } from '../webview/contract'
 import { ExperimentsOutput } from '../../cli/dvc/contract'
 import { PersistenceKey } from '../../persistence/constants'
-import { PathSelectionModel } from '../../path/selection/model'
+import { PathSelectionModel, Status } from '../../path/selection/model'
 
 export class ColumnsModel extends PathSelectionModel<Column> {
   private columnOrderState: string[] = []
@@ -91,6 +91,20 @@ export class ColumnsModel extends PathSelectionModel<Column> {
     })
   }
 
+  public getTerminalNodes(): (Column & { selected: boolean })[] {
+    return this.data
+      .filter(element => !element.hasChildren)
+      .map(element => ({ ...element, selected: !!this.status[element.path] }))
+  }
+
+  public getSelected() {
+    return (
+      this.data.filter(
+        element => this.status[element.path] !== Status.UNSELECTED
+      ) || []
+    )
+  }
+
   public hasNonDefaultColumns() {
     return this.data.length > 1
   }

extension/src/path/selection/model.ts

@@ -35,20 +35,6 @@ export abstract class PathSelectionModel<
     }
   }
 
-  public getSelected() {
-    return (
-      this.data.filter(
-        element => this.status[element.path] !== Status.UNSELECTED
-      ) || []
-    )
-  }
-
-  public getTerminalNodes(): (T & { selected: boolean })[] {
-    return this.data
-      .filter(element => !element.hasChildren)
-      .map(element => ({ ...element, selected: !!this.status[element.path] }))
-  }
-
   public toggleStatus(path: string) {
     const status = this.getNextStatus(path)
     return this.setStatus(path, status)
@@ -159,4 +145,6 @@ export abstract class PathSelectionModel<
   abstract getChildren(
     ...args: unknown[]
   ): (T & { descendantStatuses: Status[]; label: string; status: Status })[]
+
+  abstract getTerminalNodes(): (T & { selected: boolean })[]
 }

extension/src/path/selection/quickPick.test.ts

@@ -4,6 +4,7 @@ import { quickPickManyValues } from '../../vscode/quickPick'
 import { Toast } from '../../vscode/toast'
 import { Title } from '../../vscode/title'
 import { PathType } from '../../plots/paths/collect'
+import { EXPERIMENT_WORKSPACE_ID } from '../../cli/dvc/contract'
 
 jest.mock('../../vscode/quickPick')
 jest.mock('../../vscode/toast')
@@ -36,13 +37,15 @@ describe('pickPaths', () => {
 
   it('should call the quickPick with the correct items', async () => {
     mockedQuickPickManyValues.mockResolvedValueOnce([])
+    const revisions = new Set([EXPERIMENT_WORKSPACE_ID])
 
     const plotPaths = [
       {
         hasChildren: false,
         label: 'loss.tsv',
         parentPath: 'logs',
         path: join('logs', 'loss.tsv'),
+        revisions,
         selected: true,
         type: new Set([PathType.TEMPLATE_SINGLE])
       },
@@ -51,6 +54,7 @@ describe('pickPaths', () => {
         label: 'acc.tsv',
         parentPath: 'logs',
         path: join('logs', 'acc.tsv'),
+        revisions,
         selected: true,
         type: new Set([PathType.TEMPLATE_SINGLE])
       },
@@ -59,6 +63,7 @@ describe('pickPaths', () => {
         label: 'fun.tsv',
         parentPath: 'logs',
         path: join('logs', 'fun.tsv'),
+        revisions,
         selected: false,
         type: new Set([PathType.TEMPLATE_SINGLE])
       }

extension/src/plots/index.ts

@@ -121,14 +121,15 @@ export class Plots extends BaseRepository<TPlotsData> {
   }
 
   public getPathStatuses() {
-    return this.paths.getTerminalNodeStatuses()
+    return this.paths.getTerminalNodeStatuses(undefined)
   }
 
   public getScale() {
     return collectScale(this.paths.getTerminalNodes())
   }
 
   private notifyChanged() {
+    this.paths.setSelectedRevisions(this.plots.getSelectedRevisions())
     this.pathsChanged.fire()
     this.fetchMissingOrSendPlots()
   }
@@ -187,9 +188,7 @@ export class Plots extends BaseRepository<TPlotsData> {
           await this.plots.transformAndSetExperiments(data)
         }
 
-        this.plots.setComparisonOrder()
-
-        this.fetchMissingOrSendPlots()
+        this.notifyChanged()
       })
     )
   }
@@ -210,7 +209,7 @@ export class Plots extends BaseRepository<TPlotsData> {
       this.data.onDidUpdate(async ({ data, revs }) => {
         await Promise.all([
           this.plots.transformAndSetPlots(data, revs),
-          this.paths.transformAndSet(data)
+          this.paths.transformAndSet(data, revs, this.plots.getCLIIdToLabel())
         ])
         this.notifyChanged()
       })

extension/src/plots/model/index.ts

@@ -291,6 +291,10 @@ export class PlotsModel extends ModelWithPersistence {
     this.persist(PersistenceKey.PLOT_COMPARISON_ORDER, this.comparisonOrder)
   }
 
+  public getSelectedRevisions() {
+    return this.experiments.getSelectedRevisions().map(({ label }) => label)
+  }
+
   public setSelectedMetrics(selectedMetrics: string[]) {
     this.selectedMetrics = selectedMetrics
     this.setMetricOrder()
@@ -349,6 +353,16 @@ export class PlotsModel extends ModelWithPersistence {
     return this.multiSourceEncoding
   }
 
+  public getCLIIdToLabel() {
+    const mapping: { [shortSha: string]: string } = {}
+
+    for (const rev of this.experiments.getRevisions()) {
+      mapping[this.getCLIId(rev)] = rev
+    }
+
+    return mapping
+  }
+
   private removeStaleData() {
     return Promise.all([
       this.removeStaleBranches(),
@@ -397,24 +411,10 @@ export class PlotsModel extends ModelWithPersistence {
     this.fetchedRevs.delete(id)
   }
 
-  private getCLIIdToLabel() {
-    const mapping: { [shortSha: string]: string } = {}
-
-    for (const rev of this.experiments.getRevisions()) {
-      mapping[this.getCLIId(rev)] = rev
-    }
-
-    return mapping
-  }
-
   private getCLIId(label: string) {
     return this.branchRevisions[label] || label
   }
 
-  private getSelectedRevisions() {
-    return this.experiments.getSelectedRevisions().map(({ label }) => label)
-  }
-
   private getPlots(
     checkpointPlots: CheckpointPlot[],
     selectedExperiments: string[]