Add option to filter experiments to starred

extension/src/experiments/columns/like.ts

@@ -0,0 +1,22 @@
+import { Column } from '../webview/contract'
+
+export type ColumnLike = { label: string; path: string; types?: string[] }
+
+const starredColumnLike: ColumnLike = {
+  label: '$(star-full)',
+  path: 'starred',
+  types: ['boolean']
+}
+
+export const addStarredToColumns = (
+  columns: Column[] | undefined
+): ColumnLike[] | undefined => {
+  if (!columns?.length) {
+    return
+  }
+
+  return [
+    ...columns.map(({ label, path, types }) => ({ label, path, types })),
+    starredColumnLike
+  ]
+}

extension/src/experiments/columns/quickPick.test.ts

@@ -1,4 +1,4 @@
-import { pickFromColumns } from './quickPick'
+import { pickFromColumnLikes } from './quickPick'
 import { appendColumnToPath, buildMetricOrParamPath } from './paths'
 import { quickPickValue } from '../../vscode/quickPick'
 import { Toast } from '../../vscode/toast'
@@ -17,34 +17,25 @@ beforeEach(() => {
   jest.resetAllMocks()
 })
 
-describe('pickFromColumns', () => {
+describe('pickFromColumnLikes', () => {
   const params = ColumnType.PARAMS
   const paramsYaml = 'params.yaml'
   const paramsYamlPath = buildMetricOrParamPath(params, paramsYaml)
   const epochsParamPath = appendColumnToPath(paramsYamlPath, 'epochs')
   const epochsParam = {
-    hasChildren: false,
     label: 'epochs',
-    maxNumber: 5,
-    maxStringLength: 1,
-    minNumber: 2,
-    parentPath: paramsYamlPath,
     path: epochsParamPath,
-    type: ColumnType.PARAMS,
     types: ['number']
   }
 
   const paramsYamlParam = {
-    hasChildren: true,
     label: paramsYaml,
-    parentPath: params,
-    path: paramsYamlPath,
-    type: ColumnType.PARAMS
+    path: paramsYamlPath
   }
   const exampleColumns = [epochsParam, paramsYamlParam]
 
   it('should return early if no columns are provided', async () => {
-    const picked = await pickFromColumns([], {
+    const picked = await pickFromColumnLikes([], {
       title: "can't pick from no columns" as Title
     })
     expect(picked).toBeUndefined()
@@ -54,7 +45,7 @@ describe('pickFromColumns', () => {
 
   it('should invoke a QuickPick with the correct options', async () => {
     const title = 'Test title' as Title
-    await pickFromColumns(exampleColumns, { title })
+    await pickFromColumnLikes(exampleColumns, { title })
     expect(mockedQuickPickValue).toBeCalledWith(
       [
         {

extension/src/experiments/columns/quickPick.ts

@@ -1,23 +1,23 @@
+import { ColumnLike } from './like'
 import { definedAndNonEmpty } from '../../util/array'
 import {
   QuickPickOptionsWithTitle,
   quickPickValue
 } from '../../vscode/quickPick'
 import { Toast } from '../../vscode/toast'
-import { Column } from '../webview/contract'
 
-export const pickFromColumns = (
-  columns: Column[] | undefined,
+export const pickFromColumnLikes = (
+  columnLikes: ColumnLike[] | undefined,
   quickPickOptions: QuickPickOptionsWithTitle
-) => {
-  if (!definedAndNonEmpty(columns)) {
+): Thenable<ColumnLike | undefined> => {
+  if (!definedAndNonEmpty(columnLikes)) {
     return Toast.showError('There are no columns to select from.')
   }
-  return quickPickValue<Column>(
-    columns.map(column => ({
-      description: column.path,
-      label: column.label,
-      value: column
+  return quickPickValue<ColumnLike>(
+    columnLikes.map(columnLike => ({
+      description: columnLike.path,
+      label: columnLike.label,
+      value: columnLike
     })),
     quickPickOptions
   )

extension/src/experiments/index.ts

@@ -1,4 +1,5 @@
 import { Event, EventEmitter, Memento } from 'vscode'
+import { addStarredToColumns } from './columns/like'
 import { setContextForEditorTitleIcons } from './context'
 import { ExperimentsModel } from './model'
 import { pickExperiments } from './model/quickPicks'
@@ -219,10 +220,13 @@ export class Experiments extends BaseRepository<TableData> {
 
   public async addFilter() {
     const columns = this.columns.getTerminalNodes()
-    const filterToAdd = await pickFilterToAdd(columns)
+    const columnLikes = addStarredToColumns(columns)
+
+    const filterToAdd = await pickFilterToAdd(columnLikes)
     if (!filterToAdd) {
       return
     }
+
     this.experiments.addFilter(filterToAdd)
     return this.notifyChanged()
   }
@@ -335,7 +339,7 @@ export class Experiments extends BaseRepository<TableData> {
   }
 
   public getCheckpoints(id: string) {
-    return this.experiments.getCheckpoints(id)
+    return this.experiments.getCheckpointsWithType(id)
   }
 
   public sendInitialWebviewData() {

extension/src/experiments/model/filterBy/index.ts

@@ -40,7 +40,7 @@ const evaluateBoolean = (
 ): boolean =>
   typeof valueToEvaluate === 'boolean' && valueToEvaluate === filterValue
 
-const evaluate = <T>(
+const evaluate = <T extends string | number | boolean>(
   valueToEvaluate: T,
   operator: Operator,
   filterValue: T

extension/src/experiments/model/filterBy/quickPick.ts

@@ -4,8 +4,8 @@ import { getInput } from '../../../vscode/inputBox'
 import { quickPickManyValues, quickPickValue } from '../../../vscode/quickPick'
 import { Title } from '../../../vscode/title'
 import { Toast } from '../../../vscode/toast'
-import { pickFromColumns } from '../../columns/quickPick'
-import { Column } from '../../webview/contract'
+import { ColumnLike } from '../../columns/like'
+import { pickFromColumnLikes } from '../../columns/quickPick'
 
 export const operators = [
   {
@@ -84,9 +84,9 @@ const addFilterValue = async (path: string, operator: Operator) => {
 }
 
 export const pickFilterToAdd = async (
-  columns: Column[] | undefined
+  columns: ColumnLike[] | undefined
 ): Promise<FilterDefinition | undefined> => {
-  const picked = await pickFromColumns(columns, {
+  const picked = await pickFromColumnLikes(columns, {
     title: Title.SELECT_PARAM_OR_METRIC_FILTER
   })
   if (!picked) {

extension/src/experiments/model/index.ts

@@ -128,7 +128,7 @@ export class ExperimentsModel extends ModelWithPersistence {
 
   public toggleStatus(id: string) {
     if (
-      this.flattenExperiments().find(({ id: queuedId }) => queuedId === id)
+      this.getFlattenedExperiments().find(({ id: queuedId }) => queuedId === id)
         ?.queued
     ) {
       return
@@ -236,7 +236,7 @@ export class ExperimentsModel extends ModelWithPersistence {
   }
 
   public getSelectedExperiments() {
-    return this.getSelectedFromList(() => this.flattenExperiments())
+    return this.getSelectedFromList(() => this.getFlattenedExperiments())
   }
 
   public setSelected(selectedExperiments: Experiment[]) {
@@ -288,6 +288,7 @@ export class ExperimentsModel extends ModelWithPersistence {
   public getExperiments(): (ExperimentWithType & {
     hasChildren: boolean
     selected?: boolean
+    starred: boolean
   })[] {
     return [
       {
@@ -302,8 +303,8 @@ export class ExperimentsModel extends ModelWithPersistence {
           type: ExperimentType.BRANCH
         }
       }),
-      ...this.flattenExperiments().map(experiment => ({
-        ...this.addDetails(experiment),
+      ...this.getFlattenedExperiments().map(experiment => ({
+        ...experiment,
         hasChildren: definedAndNonEmpty(
           this.checkpointsByTip.get(experiment.id)
         ),
@@ -324,9 +325,7 @@ export class ExperimentsModel extends ModelWithPersistence {
 
   public getExperimentsWithCheckpoints(): ExperimentWithCheckpoints[] {
     return this.getExperiments().map(experiment => {
-      const checkpoints = this.checkpointsByTip
-        .get(experiment.id)
-        ?.map(checkpoint => this.addDetails(checkpoint))
+      const checkpoints = this.getCheckpointsWithType(experiment.id)
       if (!definedAndNonEmpty(checkpoints)) {
         return experiment
       }
@@ -350,7 +349,7 @@ export class ExperimentsModel extends ModelWithPersistence {
     return this.splitExperimentsByQueued(true)
   }
 
-  public getCheckpoints(
+  public getCheckpointsWithType(
     id: string
   ): (Experiment & { type: ExperimentType })[] | undefined {
     return this.checkpointsByTip.get(id)?.map(checkpoint => ({
@@ -388,8 +387,8 @@ export class ExperimentsModel extends ModelWithPersistence {
 
   public getExperimentCount() {
     return sum([
-      this.flattenCheckpoints().length,
-      this.flattenExperiments().length,
+      this.getFlattenedCheckpoints().length,
+      this.getFlattenedExperiments().length,
       this.branches.length,
       1
     ])
@@ -404,8 +403,8 @@ export class ExperimentsModel extends ModelWithPersistence {
     return [
       this.workspace,
       ...this.branches,
-      ...this.flattenExperiments(),
-      ...this.flattenCheckpoints()
+      ...this.getFlattenedExperiments(),
+      ...this.getFlattenedCheckpoints()
     ]
   }
 
@@ -417,13 +416,11 @@ export class ExperimentsModel extends ModelWithPersistence {
           filters
         )
         if (!checkpoints) {
-          return this.addDetails(experiment)
+          return experiment
         }
         return {
-          ...this.addDetails(experiment),
-          subRows: checkpoints.map(checkpoint => ({
-            ...this.addDetails(checkpoint)
-          }))
+          ...experiment,
+          subRows: checkpoints
         }
       })
       .filter((row: Row) => this.filterTableRow(row, filters))
@@ -445,7 +442,7 @@ export class ExperimentsModel extends ModelWithPersistence {
     const acc: ExperimentWithType[] = []
 
     for (const experiment of this.getCurrentExperiments()) {
-      const checkpoints = this.getCheckpoints(experiment.id) || []
+      const checkpoints = this.getCheckpointsWithType(experiment.id) || []
       collectFiltered(acc, this.getFilters(), experiment, checkpoints)
     }
 
@@ -456,37 +453,49 @@ export class ExperimentsModel extends ModelWithPersistence {
     sha: string,
     filters: FilterDefinition[]
   ) {
-    const checkpoints = this.checkpointsByTip.get(sha)
+    const checkpoints = this.getCheckpoints(sha)
     if (!checkpoints) {
       return
     }
     const { unfiltered } = splitExperimentsByFilters(filters, checkpoints)
     return unfiltered
   }
 
+  private getCheckpoints(id: string) {
+    return this.checkpointsByTip
+      .get(id)
+      ?.map(checkpoint => this.addDetails(checkpoint))
+  }
+
   private getExperimentsByBranch(branch: Experiment) {
-    const experiments = this.experimentsByBranch.get(branch.label)
+    const experiments = this.experimentsByBranch
+      .get(branch.label)
+      ?.map(experiment => this.addDetails(experiment))
     if (!experiments) {
       return
     }
     return sortExperiments(this.getSorts(), experiments)
   }
 
-  private flattenExperiments() {
-    return flattenMapValues(this.experimentsByBranch)
+  private getFlattenedExperiments() {
+    return flattenMapValues(this.experimentsByBranch).map(experiment =>
+      this.addDetails(experiment)
+    )
   }
 
   private splitExperimentsByQueued(getQueued = false) {
-    return this.flattenExperiments().filter(({ queued }) => {
+    return this.getFlattenedExperiments().filter(({ queued }) => {
       if (getQueued) {
         return queued
       }
       return !queued
     })
   }
 
-  private flattenCheckpoints() {
-    return flattenMapValues(this.checkpointsByTip)
+  private getFlattenedCheckpoints() {
+    return flattenMapValues(this.checkpointsByTip).map(checkpoint =>
+      this.addDetails(checkpoint)
+    )
   }
 
   private setColoredStatus() {
@@ -546,38 +555,28 @@ export class ExperimentsModel extends ModelWithPersistence {
     return this.persist(PersistenceKey.EXPERIMENTS_STATUS, this.coloredStatus)
   }
 
-  private addStarred(experiment: Experiment) {
+  private addDetails(experiment: Experiment) {
     const { id } = experiment
 
-    if (!this.isStarred(id)) {
-      return experiment
-    }
-
-    return {
-      ...experiment,
-      starred: true
-    }
-  }
+    const starred = !!this.isStarred(id)
 
-  private addSelected(experiment: Experiment) {
-    const { id } = experiment
     if (!hasKey(this.coloredStatus, id)) {
-      return experiment
+      return {
+        ...experiment,
+        starred
+      }
     }
 
     const selected = this.isSelected(id)
 
     return {
       ...experiment,
       displayColor: this.getDisplayColor(id),
-      selected
+      selected,
+      starred
     }
   }
 
-  private addDetails(experiment: Experiment) {
-    return this.addStarred(this.addSelected(experiment))
-  }
-
   private getDisplayColor(id: string) {
     const color = this.coloredStatus[id]
     if (!color) {

extension/src/experiments/model/sortBy/quickPick.ts

@@ -3,11 +3,11 @@ import { definedAndNonEmpty } from '../../../util/array'
 import { quickPickManyValues, quickPickValue } from '../../../vscode/quickPick'
 import { Title } from '../../../vscode/title'
 import { Toast } from '../../../vscode/toast'
-import { pickFromColumns } from '../../columns/quickPick'
+import { pickFromColumnLikes } from '../../columns/quickPick'
 import { Column } from '../../webview/contract'
 
 export const pickSortToAdd = async (columns: Column[]) => {
-  const picked = await pickFromColumns(columns, {
+  const picked = await pickFromColumnLikes(columns, {
     title: Title.SELECT_PARAM_OR_METRIC_SORT
   })
   if (picked === undefined) {