Merge pull request #1950 from iterative/star-experiments

Star experiments in table

Author: wolmir

extension/src/experiments/index.ts

@@ -176,6 +176,11 @@ export class Experiments extends BaseRepository<TableData> {
     return this.columns.getTerminalNodeStatuses()
   }
 
+  public toggleExperimentStars(ids: string[]) {
+    this.experiments.toggleStars(ids)
+    this.notifyChanged()
+  }
+
   public toggleExperimentStatus(id: string) {
     const selected = this.experiments.isSelected(id)
     if (!selected && !this.experiments.canSelect()) {
@@ -532,6 +537,8 @@ export class Experiments extends BaseRepository<TableData> {
               RegisteredCommands.EXPERIMENT_TOGGLE,
               { dvcRoot: this.dvcRoot, id: message.payload }
             )
+          case MessageFromWebviewType.TOGGLE_EXPERIMENT_STAR:
+            return this.setExperimentStars(message.payload)
           case MessageFromWebviewType.HIDE_EXPERIMENTS_TABLE_COLUMN:
             return this.hideTableColumn(message.payload)
           case MessageFromWebviewType.OPEN_PARAMS_FILE_TO_THE_SIDE:
@@ -604,6 +611,15 @@ export class Experiments extends BaseRepository<TableData> {
     )
   }
 
+  private setExperimentStars(ids: string[]) {
+    this.toggleExperimentStars(ids)
+    sendTelemetryEvent(
+      EventName.VIEWS_EXPERIMENTS_TABLE_EXPERIMENT_STARS_TOGGLE,
+      undefined,
+      undefined
+    )
+  }
+
   private setColumnsStatus() {
     this.selectColumns()
     sendTelemetryEvent(

extension/src/experiments/model/index.ts

@@ -32,6 +32,8 @@ import { ModelWithPersistence } from '../../persistence/model'
 import { PersistenceKey } from '../../persistence/constants'
 import { sum } from '../../util/math'
 
+export type StarredExperiments = Record<string, boolean | undefined>
+
 type SelectedExperimentWithColor = Experiment & {
   displayColor: Color
   selected: true
@@ -56,6 +58,7 @@ export class ExperimentsModel extends ModelWithPersistence {
   private checkpointsByTip: Map<string, Experiment[]> = new Map()
   private availableColors: Color[]
   private coloredStatus: ColoredStatus
+  private starredExperiments: StarredExperiments
 
   private filters: Map<string, FilterDefinition> = new Map()
   private useFiltersForSelection = false
@@ -81,6 +84,10 @@ export class ExperimentsModel extends ModelWithPersistence {
       PersistenceKey.EXPERIMENTS_STATUS,
       {}
     )
+    this.starredExperiments = this.revive<StarredExperiments>(
+      PersistenceKey.EXPERIMENTS_STARS,
+      {}
+    )
 
     const assignedColors = new Set(
       Object.values(this.coloredStatus).filter(Boolean)
@@ -108,6 +115,13 @@ export class ExperimentsModel extends ModelWithPersistence {
     this.setColoredStatus()
   }
 
+  public toggleStars(ids: string[]) {
+    for (const id of ids) {
+      this.starredExperiments[id] = !this.starredExperiments[id]
+      this.persistStars()
+    }
+  }
+
   public toggleStatus(id: string) {
     if (
       this.flattenExperiments().find(({ id: queuedId }) => queuedId === id)
@@ -273,19 +287,19 @@ export class ExperimentsModel extends ModelWithPersistence {
   })[] {
     return [
       {
-        ...this.addSelected(this.workspace),
+        ...this.addDetails(this.workspace),
         hasChildren: false,
         type: ExperimentType.WORKSPACE
       },
       ...this.branches.map(branch => {
         return {
-          ...this.addSelected(branch),
+          ...this.addDetails(branch),
           hasChildren: false,
           type: ExperimentType.BRANCH
         }
       }),
       ...this.flattenExperiments().map(experiment => ({
-        ...this.addSelected(experiment),
+        ...this.addDetails(experiment),
         hasChildren: definedAndNonEmpty(
           this.checkpointsByTip.get(experiment.id)
         ),
@@ -300,7 +314,7 @@ export class ExperimentsModel extends ModelWithPersistence {
     return this.getExperiments().map(experiment => {
       const checkpoints = this.checkpointsByTip
         .get(experiment.id)
-        ?.map(checkpoint => this.addSelected(checkpoint))
+        ?.map(checkpoint => this.addDetails(checkpoint))
       if (!definedAndNonEmpty(checkpoints)) {
         return experiment
       }
@@ -328,24 +342,24 @@ export class ExperimentsModel extends ModelWithPersistence {
     id: string
   ): (Experiment & { type: ExperimentType })[] | undefined {
     return this.checkpointsByTip.get(id)?.map(checkpoint => ({
-      ...this.addSelected(checkpoint),
+      ...this.addDetails(checkpoint),
       type: ExperimentType.CHECKPOINT
     }))
   }
 
   public getRowData() {
     return [
-      this.addSelected(this.workspace),
+      this.addDetails(this.workspace),
       ...this.branches.map(branch => {
         const experiments = this.getExperimentsByBranch(branch)
-        const branchWithSelected = this.addSelected(branch)
+        const branchWithSelectedAndStarred = this.addDetails(branch)
 
         if (!definedAndNonEmpty(experiments)) {
-          return branchWithSelected
+          return branchWithSelectedAndStarred
         }
 
         return {
-          ...branchWithSelected,
+          ...branchWithSelectedAndStarred,
           subRows: this.getSubRows(experiments)
         }
       })
@@ -356,6 +370,10 @@ export class ExperimentsModel extends ModelWithPersistence {
     return !!this.coloredStatus[id]
   }
 
+  public isStarred(id: string) {
+    return !!this.starredExperiments[id]
+  }
+
   public getExperimentCount() {
     return sum([
       this.flattenCheckpoints().length,
@@ -387,12 +405,12 @@ export class ExperimentsModel extends ModelWithPersistence {
           filters
         )
         if (!checkpoints) {
-          return this.addSelected(experiment)
+          return this.addDetails(experiment)
         }
         return {
-          ...this.addSelected(experiment),
+          ...this.addDetails(experiment),
           subRows: checkpoints.map(checkpoint => ({
-            ...this.addSelected(checkpoint)
+            ...this.addDetails(checkpoint)
           }))
         }
       })
@@ -505,10 +523,30 @@ export class ExperimentsModel extends ModelWithPersistence {
     return this.persist(PersistenceKey.EXPERIMENTS_FILTER_BY, [...this.filters])
   }
 
+  private persistStars() {
+    return this.persist(
+      PersistenceKey.EXPERIMENTS_STARS,
+      this.starredExperiments
+    )
+  }
+
   private persistStatus() {
     return this.persist(PersistenceKey.EXPERIMENTS_STATUS, this.coloredStatus)
   }
 
+  private addStarred(experiment: Experiment) {
+    const { id } = experiment
+
+    if (!this.isStarred(id)) {
+      return experiment
+    }
+
+    return {
+      ...experiment,
+      starred: true
+    }
+  }
+
   private addSelected(experiment: Experiment) {
     const { id } = experiment
     if (!hasKey(this.coloredStatus, id)) {
@@ -524,6 +562,10 @@ export class ExperimentsModel extends ModelWithPersistence {
     }
   }
 
+  private addDetails(experiment: Experiment) {
+    return this.addStarred(this.addSelected(experiment))
+  }
+
   private getDisplayColor(id: string) {
     const color = this.coloredStatus[id]
     if (!color) {

extension/src/experiments/webview/contract.ts

@@ -21,6 +21,7 @@ export interface Experiment extends BaseExperimentFields {
   outs?: MetricOrParamColumns
   displayColor?: string
   selected?: boolean
+  starred?: boolean
   mutable?: boolean
   sha?: string
 }

extension/src/persistence/constants.ts

@@ -2,6 +2,7 @@ export enum PersistenceKey {
   EXPERIMENTS_FILTER_BY = 'experimentsFilterBy:',
   EXPERIMENTS_SORT_BY = 'experimentsSortBy:',
   EXPERIMENTS_STATUS = 'experimentsStatus:',
+  EXPERIMENTS_STARS = 'experimentsStars:',
   METRICS_AND_PARAMS_COLUMN_ORDER = 'columnsColumnOrder:',
   METRICS_AND_PARAMS_COLUMN_WIDTHS = 'columnsColumnWidths:',
   METRICS_AND_PARAMS_STATUS = 'columnsStatus:',

extension/src/telemetry/constants.ts

@@ -32,6 +32,8 @@ export const EventName = Object.assign(
     VIEWS_EXPERIMENTS_TABLE_COLUMNS_REORDERED:
       'views.experimentsTable.columnsReordered',
     VIEWS_EXPERIMENTS_TABLE_CREATED: 'views.experimentsTable.created',
+    VIEWS_EXPERIMENTS_TABLE_EXPERIMENT_STARS_TOGGLE:
+      'views.experimentTable.toggleStars',
     VIEWS_EXPERIMENTS_TABLE_EXPERIMENT_TOGGLE:
       'views.experimentTable.toggleStatus',
     VIEWS_EXPERIMENTS_TABLE_FOCUS_CHANGED:
@@ -189,6 +191,7 @@ export interface IEventNamePropertyMapping {
   [EventName.VIEWS_EXPERIMENTS_SORT_BY_TREE_OPENED]: DvcRootCount
   [EventName.VIEWS_EXPERIMENTS_TREE_OPENED]: DvcRootCount
   [EventName.VIEWS_EXPERIMENTS_TABLE_EXPERIMENT_TOGGLE]: undefined
+  [EventName.VIEWS_EXPERIMENTS_TABLE_EXPERIMENT_STARS_TOGGLE]: undefined
   [EventName.VIEWS_EXPERIMENTS_TABLE_CLOSED]: undefined
   [EventName.VIEWS_EXPERIMENTS_TABLE_COLUMNS_REORDERED]: undefined
   [EventName.VIEWS_EXPERIMENTS_TABLE_FOCUS_FILTERS_TREE]: undefined

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

@@ -788,6 +788,42 @@ suite('Experiments Test Suite', () => {
         undefined
       )
     })
+
+    it("should be able to handle a message to toggle an experiment's star status", async () => {
+      const { experiments, experimentsModel } =
+        setupExperimentsAndMockCommands()
+
+      const experimentsToToggle = ['exp-e7a67']
+
+      const areExperimentsStarred = (expIds: string[]): boolean =>
+        expIds
+          .map(expId =>
+            experimentsModel.getExperiments().find(({ id }) => id === expId)
+          )
+          .every(exp => exp?.starred)
+
+      expect(
+        areExperimentsStarred(experimentsToToggle),
+        'experiments are starred'
+      ).to.be.false
+
+      const webview = await experiments.showWebview()
+      const mockMessageReceived = getMessageReceivedEmitter(webview)
+      const toggleSpy = spy(experimentsModel, 'toggleStars')
+
+      mockMessageReceived.fire({
+        payload: experimentsToToggle,
+        type: MessageFromWebviewType.TOGGLE_EXPERIMENT_STAR
+      })
+
+      expect(toggleSpy).to.be.calledWith(experimentsToToggle)
+      toggleSpy.resetHistory()
+
+      expect(
+        areExperimentsStarred(experimentsToToggle),
+        'experiments have been starred'
+      ).to.be.true
+    })
   })
 
   describe('Sorting', () => {

extension/src/webview/contract.ts

@@ -29,6 +29,7 @@ export enum MessageFromWebviewType {
   RESIZE_PLOTS = 'resize-plots',
   SORT_COLUMN = 'sort-column',
   TOGGLE_EXPERIMENT = 'toggle-experiment',
+  TOGGLE_EXPERIMENT_STAR = 'toggle-experiment-star',
   HIDE_EXPERIMENTS_TABLE_COLUMN = 'hide-experiments-table-column',
   SELECT_EXPERIMENTS = 'select-experiments',
   SELECT_COLUMNS = 'select-columns',
@@ -67,6 +68,10 @@ export type MessageFromWebview =
       type: MessageFromWebviewType.TOGGLE_EXPERIMENT
       payload: string
     }
+  | {
+      type: MessageFromWebviewType.TOGGLE_EXPERIMENT_STAR
+      payload: string[]
+    }
   | {
       type: MessageFromWebviewType.HIDE_EXPERIMENTS_TABLE_COLUMN
       payload: string