Add select first experiment table columns quick pick (#4294)

mattseddon · web-flow · commit 634f7e96edac · 2023-07-18T21:24:28.000Z
diff --git a/extension/package.json b/extension/package.json
@@ -571,6 +571,12 @@
         "category": "DVC",
         "icon": "$(list-filter)"
       },
+      {
+        "title": "Select Columns to Display First in the Experiments Table",
+        "command": "dvc.views.experimentsColumnsTree.selectFirstColumns",
+        "category": "DVC",
+        "icon": "$(arrow-left)"
+      },
       {
         "title": "Select Plots to Display",
         "command": "dvc.views.plotsPathsTree.selectPlots",
@@ -919,6 +925,10 @@
           "command": "dvc.views.experimentsColumnsTree.selectColumns",
           "when": "dvc.commands.available && dvc.project.available"
         },
+        {
+          "command": "dvc.views.experimentsColumnsTree.selectFirstColumns",
+          "when": "dvc.commands.available && dvc.project.available"
+        },
         {
           "command": "dvc.views.experiments.applyExperiment",
           "when": "false"
@@ -1284,6 +1294,11 @@
           "when": "view == dvc.views.experimentsColumnsTree",
           "group": "navigation@2"
         },
+        {
+          "command": "dvc.views.experimentsColumnsTree.selectFirstColumns",
+          "when": "view == dvc.views.experimentsColumnsTree",
+          "group": "navigation@3"
+        },
         {
           "command": "dvc.stopAllRunningExperiments",
           "when": "view == dvc.views.experimentsTree && dvc.experiments.webview.active && dvc.experiment.running",
diff --git a/extension/src/commands/external.ts b/extension/src/commands/external.ts
@@ -51,6 +51,7 @@ export enum RegisteredCliCommands {
 
 export enum RegisteredCommands {
   EXPERIMENT_COLUMNS_SELECT = 'dvc.views.experimentsColumnsTree.selectColumns',
+  EXPERIMENT_COLUMNS_SELECT_FIRST = 'dvc.views.experimentsColumnsTree.selectFirstColumns',
   EXPERIMENT_FILTER_ADD = 'dvc.addExperimentsTableFilter',
   EXPERIMENT_FILTER_ADD_STARRED = 'dvc.addStarredExperimentsTableFilter',
   EXPERIMENT_FILTER_REMOVE = 'dvc.views.experimentsFilterByTree.removeFilter',
diff --git a/extension/src/experiments/columns/model.ts b/extension/src/experiments/columns/model.ts
@@ -94,6 +94,17 @@ export class ColumnsModel extends PathSelectionModel<Column> {
     this.statusChanged?.fire()
   }
 
+  public selectFirst(firstColumns: string[]) {
+    const columnOrder = [
+      'id',
+      ...firstColumns,
+      ...this.getColumnOrder().filter(
+        column => !['id', ...firstColumns].includes(column)
+      )
+    ]
+    this.setColumnOrder(columnOrder)
+  }
+
   public setColumnWidth(id: string, width: number) {
     this.columnWidthsState[id] = width
     this.persist(
diff --git a/extension/src/experiments/commands/register.ts b/extension/src/experiments/commands/register.ts
@@ -190,6 +190,12 @@ const registerExperimentQuickPickCommands = (
       experiments.selectColumns(getDvcRootFromContext(context))
   )
 
+  internalCommands.registerExternalCommand(
+    RegisteredCommands.EXPERIMENT_COLUMNS_SELECT_FIRST,
+    (context: Context) =>
+      experiments.selectFirstColumns(getDvcRootFromContext(context))
+  )
+
   internalCommands.registerExternalCommand(
     RegisteredCommands.EXPERIMENT_STOP,
     () => experiments.selectExperimentsToStop()
diff --git a/extension/src/experiments/index.ts b/extension/src/experiments/index.ts
@@ -45,6 +45,7 @@ import { checkSignalFile, pollSignalFileForProcess } from '../fileSystem'
 import { DVCLIVE_ONLY_RUNNING_SIGNAL_FILE } from '../cli/dvc/constants'
 import { Response } from '../vscode/response'
 import { Pipeline } from '../pipeline'
+import { definedAndNonEmpty } from '../util/array'
 
 export const ExperimentsScale = {
   ...omit(ColumnType, 'TIMESTAMP'),
@@ -356,7 +357,7 @@ export class Experiments extends BaseRepository<TableData> {
   public async selectColumns() {
     const columns = this.columns.getTerminalNodes()
 
-    const selected = await pickPaths('columns', columns)
+    const selected = await pickPaths(columns, Title.SELECT_COLUMNS)
     if (!selected) {
       return
     }
@@ -365,6 +366,21 @@ export class Experiments extends BaseRepository<TableData> {
     return this.notifyChanged()
   }
 
+  public async selectFirstColumns() {
+    const columns = this.columns.getTerminalNodes()
+
+    const selected = await pickPaths(
+      columns.map(column => ({ ...column, selected: false })),
+      Title.SELECT_FIRST_COLUMNS
+    )
+    if (!definedAndNonEmpty(selected)) {
+      return
+    }
+
+    this.columns.selectFirst(selected.map(({ path }) => path))
+    return this.notifyChanged()
+  }
+
   public pickCommitOrExperiment() {
     return pickExperiment(
       this.experiments.getCommitsAndExperiments(),
diff --git a/extension/src/experiments/workspace.ts b/extension/src/experiments/workspace.ts
@@ -101,6 +101,10 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
     return this.getRepositoryThenUpdate('selectColumns', overrideRoot)
   }
 
+  public selectFirstColumns(overrideRoot?: string) {
+    return this.getRepositoryThenUpdate('selectFirstColumns', overrideRoot)
+  }
+
   public async selectExperimentsToStop() {
     const cwd = await this.getFocusedOrOnlyOrPickProject()
     if (!cwd) {
@@ -407,7 +411,8 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
       | 'addStarredSort'
       | 'removeSorts'
       | 'selectExperimentsToPlot'
-      | 'selectColumns',
+      | 'selectColumns'
+      | 'selectFirstColumns',
     overrideRoot?: string
   ) {
     const dvcRoot = await this.getDvcRoot(overrideRoot)
diff --git a/extension/src/path/selection/quickPick.test.ts b/extension/src/path/selection/quickPick.test.ts
@@ -21,15 +21,15 @@ beforeEach(() => {
 describe('pickPaths', () => {
   it('should not call quickPickManyValues if undefined is provided', async () => {
     mockedQuickPickManyValues.mockResolvedValueOnce([])
-    await pickPaths('plots', undefined)
+    await pickPaths(undefined, Title.SELECT_COLUMNS)
 
     expect(mockedShowError).toHaveBeenCalledTimes(1)
     expect(mockedQuickPickManyValues).not.toHaveBeenCalled()
   })
 
   it('should not call quickPickManyValues if no plots paths are provided', async () => {
     mockedQuickPickManyValues.mockResolvedValueOnce([])
-    await pickPaths('plots', [])
+    await pickPaths([], Title.SELECT_COLUMNS)
 
     expect(mockedShowError).toHaveBeenCalledTimes(1)
     expect(mockedQuickPickManyValues).not.toHaveBeenCalled()
@@ -69,7 +69,7 @@ describe('pickPaths', () => {
       }
     ]
 
-    await pickPaths('plots', plotPaths)
+    await pickPaths(plotPaths, Title.SELECT_PLOTS)
 
     expect(mockedShowError).not.toHaveBeenCalled()
     expect(mockedQuickPickManyValues).toHaveBeenCalledTimes(1)
diff --git a/extension/src/path/selection/quickPick.ts b/extension/src/path/selection/quickPick.ts
@@ -33,10 +33,13 @@ const collectItems = <T extends PathType>(
 }
 
 export const pickPaths = <T extends PathType>(
-  type: 'plots' | 'columns',
-  paths?: PathWithSelected<T>[]
+  paths: PathWithSelected<T>[] | undefined,
+  title:
+    | typeof Title.SELECT_PLOTS
+    | typeof Title.SELECT_COLUMNS
+    | typeof Title.SELECT_FIRST_COLUMNS
 ): Thenable<PathWithSelected<T>[] | undefined> => {
-  const title = type === 'plots' ? Title.SELECT_PLOTS : Title.SELECT_COLUMNS
+  const type = Title.SELECT_PLOTS === title ? 'plots' : 'columns'
 
   if (!definedAndNonEmpty(paths)) {
     return Toast.showError(`There are no ${type} to select.`)
diff --git a/extension/src/plots/index.ts b/extension/src/plots/index.ts
@@ -96,7 +96,7 @@ export class Plots extends BaseRepository<TPlotsData> {
   public async selectPlots() {
     const paths = this.paths.getTerminalNodes()
 
-    const selected = await pickPaths('plots', paths)
+    const selected = await pickPaths(paths, Title.SELECT_PLOTS)
     if (!selected) {
       return
     }
diff --git a/extension/src/telemetry/constants.ts b/extension/src/telemetry/constants.ts
@@ -131,6 +131,7 @@ export interface IEventNamePropertyMapping {
   [EventName.EXPERIMENT_APPLY]: undefined
   [EventName.EXPERIMENT_BRANCH]: undefined
   [EventName.EXPERIMENT_COLUMNS_SELECT]: undefined
+  [EventName.EXPERIMENT_COLUMNS_SELECT_FIRST]: undefined
   [EventName.EXPERIMENT_FILTER_ADD]: undefined
   [EventName.EXPERIMENT_FILTER_ADD_STARRED]: undefined
   [EventName.EXPERIMENT_FILTER_REMOVE]: undefined
diff --git a/extension/src/test/suite/experiments/columns/tree.test.ts b/extension/src/test/suite/experiments/columns/tree.test.ts
@@ -1,7 +1,7 @@
 import { afterEach, beforeEach, describe, it, suite } from 'mocha'
 import { expect } from 'chai'
-import { stub, restore } from 'sinon'
-import { commands } from 'vscode'
+import { stub, restore, SinonStub } from 'sinon'
+import { QuickPickItem, commands, window } from 'vscode'
 import { Disposable } from '../../../../extension'
 import { WorkspaceExperiments } from '../../../../experiments/workspace'
 import { dvcDemoPath } from '../../../util'
@@ -10,9 +10,13 @@ import {
   appendColumnToPath,
   buildMetricOrParamPath
 } from '../../../../experiments/columns/paths'
-import { buildExperiments } from '../util'
+import { buildExperiments, stubWorkspaceExperimentsGetters } from '../util'
 import { Status } from '../../../../path/selection/model'
 import { ColumnType } from '../../../../experiments/webview/contract'
+import {
+  QuickPickItemWithValue,
+  QuickPickOptionsWithTitle
+} from '../../../../vscode/quickPick'
 
 suite('Experiments Columns Tree Test Suite', () => {
   const paramsFile = 'params.yaml'
@@ -349,5 +353,67 @@ suite('Experiments Columns Tree Test Suite', () => {
         'the grandparent is now unselected'
       ).to.equal(Status.UNSELECTED)
     })
+
+    it('should be able to display selected columns first with dvc.views.experimentsColumnsTree.selectFirstColumns', async () => {
+      const { experiments, columnsModel } =
+        stubWorkspaceExperimentsGetters(disposable)
+      await experiments.isReady()
+
+      const columnsOrder = columnsModel.getColumnOrder()
+
+      const firstColumns = []
+      const otherColumns = []
+      for (const column of columnsOrder) {
+        if (column === 'id') {
+          continue
+        }
+        if (
+          [
+            'params:params.yaml:learning_rate',
+            'params:params.yaml:dvc_logs_dir'
+          ].includes(column)
+        ) {
+          firstColumns.push(column)
+          continue
+        }
+        otherColumns.push(column)
+      }
+
+      ;(
+        stub(window, 'showQuickPick') as SinonStub<
+          [items: readonly QuickPickItem[], options: QuickPickOptionsWithTitle],
+          Thenable<QuickPickItemWithValue<{ path: string }>[] | undefined>
+        >
+      ).resolves(
+        firstColumns.map(
+          path =>
+            ({
+              label: path,
+              value: { path }
+            }) as QuickPickItemWithValue<{ path: string }>
+        )
+      )
+
+      const orderUpdated = new Promise(resolve =>
+        disposable.track(
+          experiments.onDidChangeColumnOrderOrStatus(() => {
+            resolve(undefined)
+          })
+        )
+      )
+
+      await Promise.all([
+        commands.executeCommand(
+          RegisteredCommands.EXPERIMENT_COLUMNS_SELECT_FIRST
+        ),
+        orderUpdated
+      ])
+
+      expect(columnsModel.getColumnOrder()).to.deep.equal([
+        'id',
+        ...firstColumns,
+        ...otherColumns
+      ])
+    })
   })
 })
diff --git a/extension/src/test/suite/experiments/util.ts b/extension/src/test/suite/experiments/util.ts
@@ -255,8 +255,14 @@ export const stubWorkspaceExperimentsGetters = (
   disposer: Disposer,
   dvcRoot = dvcDemoPath
 ) => {
-  const { dvcExecutor, dvcRunner, experiments, experimentsModel, messageSpy } =
-    buildExperiments({ disposer })
+  const {
+    columnsModel,
+    dvcExecutor,
+    dvcRunner,
+    experiments,
+    experimentsModel,
+    messageSpy
+  } = buildExperiments({ disposer })
 
   const mockGetOnlyOrPickProject = stub(
     WorkspaceExperiments.prototype,
@@ -269,6 +275,7 @@ export const stubWorkspaceExperimentsGetters = (
   ).returns(experiments)
 
   return {
+    columnsModel,
     dvcExecutor,
     dvcRunner,
     experiments,
diff --git a/extension/src/vscode/title.ts b/extension/src/vscode/title.ts
@@ -23,6 +23,7 @@ export enum Title {
   SELECT_EXPERIMENTS_REMOVE = 'Select Experiment(s) to Remove',
   SELECT_EXPERIMENTS_TO_PLOT = 'Select up to 7 Experiments to Display in Plots',
   SELECT_FILTERS_TO_REMOVE = 'Select Filter(s) to Remove',
+  SELECT_FIRST_COLUMNS = 'Select Column(s) to Display First in the Experiments Table',
   SELECT_FOCUSED_PROJECTS = 'Select Project(s) to Focus (set dvc.focusedProjects)',
   SELECT_OPERATOR = 'Select an Operator',
   SELECT_PARAM_OR_METRIC_FILTER = 'Select a Param or Metric to Filter by',
diff --git a/webview/src/experiments/components/Experiments.tsx b/webview/src/experiments/components/Experiments.tsx
@@ -122,7 +122,7 @@ const defaultColumn: Partial<ColumnDef<Commit>> = {
 export const ExperimentsTable: React.FC = () => {
   const {
     columns: columnsData,
-    columnOrder: initialColumnOrder,
+    columnOrder: columnOrderData,
     columnWidths,
     hasColumns,
     hasConfig,
@@ -134,7 +134,7 @@ export const ExperimentsTable: React.FC = () => {
   const [columns, setColumns] = useState(getColumns(columnsData))
   const [columnSizing, setColumnSizing] =
     useState<ColumnSizingState>(columnWidths)
-  const [columnOrder, setColumnOrder] = useState(initialColumnOrder)
+  const [columnOrder, setColumnOrder] = useState(columnOrderData)
   const resizeTimeout = useRef(0)
 
   useEffect(() => {
@@ -145,6 +145,10 @@ export const ExperimentsTable: React.FC = () => {
     setColumns(getColumns(columnsData))
   }, [columnsData])
 
+  useEffect(() => {
+    setColumnOrder(columnOrderData)
+  }, [columnOrderData])
+
   const getRowId = useCallback(
     (experiment: Commit, relativeIndex: number, parent?: TableRow<Commit>) =>
       parent ? [parent.id, experiment.id].join('.') : String(relativeIndex),

Original file line number	Diff line number	Diff line change
`@@ -96,7 +96,7 @@ export class Plots extends BaseRepository<TPlotsData> {`
`96`	`96`	`public async selectPlots() {`
`97`	`97`	`const paths = this.paths.getTerminalNodes()`
`98`	`98`
`99`		`- const selected = await pickPaths('plots', paths)`
	`99`	`+ const selected = await pickPaths(paths, Title.SELECT_PLOTS)`
`100`	`100`	`if (!selected) {`
`101`	`101`	`return`
`102`	`102`	`}`