Add experiment icons to editor/title when params file is open (#1740)

* add prototype for showing experiment icons in editor/title when params file is open

* add variations for all commands

* await result of async function

* move listener into private method

* test and refactor

* add workspace focus when params file is open

Author: mattseddon

extension/package.json

@@ -835,6 +835,11 @@
         }
       ],
       "editor/title": [
+        {
+          "command": "dvc.runExperiment",
+          "group": "navigation@0",
+          "when": "dvc.params.fileActive && !dvc.runner.running && dvc.commands.available && !dvc.experiment.checkpoints"
+        },
         {
           "command": "dvc.runExperiment",
           "group": "navigation@0",
@@ -850,20 +855,40 @@
           "group": "navigation@1",
           "when": "dvc.experiments.webviewActive && !dvc.runner.running && dvc.commands.available && dvc.experiment.checkpoints"
         },
+        {
+          "command": "dvc.resumeCheckpointExperiment",
+          "group": "navigation@1",
+          "when": "dvc.params.fileActive && !dvc.runner.running && dvc.commands.available && dvc.experiment.checkpoints"
+        },
         {
           "command": "dvc.resetAndRunCheckpointExperiment",
           "group": "navigation@2",
           "when": "dvc.experiments.webviewActive && !dvc.runner.running && dvc.commands.available && dvc.experiment.checkpoints"
         },
+        {
+          "command": "dvc.resetAndRunCheckpointExperiment",
+          "group": "navigation@2",
+          "when": "dvc.params.fileActive && !dvc.runner.running && dvc.commands.available && dvc.experiment.checkpoints"
+        },
         {
           "command": "dvc.runQueuedExperiments",
           "group": "navigation@3",
           "when": "dvc.experiments.webviewActive && !dvc.runner.running && dvc.commands.available"
         },
+        {
+          "command": "dvc.runQueuedExperiments",
+          "group": "navigation@3",
+          "when": "dvc.params.fileActive && !dvc.runner.running && dvc.commands.available"
+        },
         {
           "command": "dvc.queueExperiment",
           "group": "navigation@4",
           "when": "dvc.experiments.webviewActive && !dvc.runner.running && dvc.commands.available"
+        },
+        {
+          "command": "dvc.queueExperiment",
+          "group": "navigation@4",
+          "when": "dvc.params.fileActive && !dvc.runner.running && dvc.commands.available"
         }
       ],
       "view/item/context": [

extension/src/experiments/columns/collect.ts

@@ -1,3 +1,4 @@
+import { join } from 'path'
 import get from 'lodash/get'
 import { ValueWalkMeta, walkRepo } from './walk'
 import { joinColumnPath, METRIC_PARAM_SEPARATOR } from './paths'
@@ -10,6 +11,7 @@ import {
   ValueTree,
   ValueTreeOrError
 } from '../../cli/reader'
+import { standardizePath } from '../../fileSystem/path'
 
 const getValueType = (value: Value) => {
   if (value === null) {
@@ -205,6 +207,16 @@ const collectColumnsChanges = (
   }
 }
 
+export const collectParamsFiles = (
+  dvcRoot: string,
+  data: ExperimentsOutput
+): Set<string> => {
+  const files = Object.keys(data.workspace.baseline.data?.params || {})
+    .filter(Boolean)
+    .map(file => standardizePath(join(dvcRoot, file))) as string[]
+  return new Set(files)
+}
+
 const getData = (value: { baseline: ExperimentFieldsOrError }) =>
   value.baseline.data || {}
 

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

@@ -14,15 +14,15 @@ import {
 describe('ColumnsModel', () => {
   const exampleDvcRoot = 'test'
 
-  it('should return columns that equal the columns fixture when given the output fixture', () => {
+  it('should return columns that equal the columns fixture when given the output fixture', async () => {
     const model = new ColumnsModel('', buildMockMemento())
-    model.transformAndSet(outputFixture)
+    await model.transformAndSet(outputFixture)
     expect(model.getSelected()).toStrictEqual(columnsFixture)
   })
 
-  it('should return data that equal the deeply nested output fixture', () => {
+  it('should return data that equal the deeply nested output fixture', async () => {
     const model = new ColumnsModel('', buildMockMemento())
-    model.transformAndSet(deeplyNestedOutput)
+    await model.transformAndSet(deeplyNestedOutput)
     expect(model.getSelected()).toStrictEqual(deeplyNestedColumns)
   })
   describe('persistence', () => {

extension/src/experiments/columns/model.ts

@@ -1,5 +1,5 @@
 import { Memento } from 'vscode'
-import { collectChanges, collectColumns } from './collect'
+import { collectChanges, collectColumns, collectParamsFiles } from './collect'
 import { splitColumnPath } from './paths'
 import { Column, ColumnType } from '../webview/contract'
 import { ExperimentsOutput } from '../../cli/reader'
@@ -10,6 +10,7 @@ export class ColumnsModel extends PathSelectionModel<Column> {
   private columnOrderState: string[] = []
   private columnWidthsState: Record<string, number> = {}
   private columnsChanges: string[] = []
+  private paramsFiles = new Set<string>()
 
   constructor(dvcRoot: string, workspaceState: Memento) {
     super(
@@ -37,6 +38,10 @@ export class ColumnsModel extends PathSelectionModel<Column> {
     return this.columnWidthsState
   }
 
+  public getParamsFiles() {
+    return this.paramsFiles
+  }
+
   public transformAndSet(data: ExperimentsOutput) {
     return Promise.all([
       this.transformAndSetColumns(data),
@@ -72,12 +77,16 @@ export class ColumnsModel extends PathSelectionModel<Column> {
     )
   }
 
-  private transformAndSetColumns(data: ExperimentsOutput) {
-    const columns = collectColumns(data)
+  private async transformAndSetColumns(data: ExperimentsOutput) {
+    const [columns, paramsFiles] = await Promise.all([
+      collectColumns(data),
+      collectParamsFiles(this.dvcRoot, data)
+    ])
 
     this.setNewStatuses(columns)
 
     this.data = columns
+    this.paramsFiles = paramsFiles
   }
 
   private transformAndSetChanges(data: ExperimentsOutput) {

extension/src/experiments/index.ts

@@ -1,4 +1,4 @@
-import { Event, EventEmitter, Memento } from 'vscode'
+import { Event, EventEmitter, Memento, window } from 'vscode'
 import { ExperimentsModel } from './model'
 import { pickExperiments } from './model/quickPicks'
 import { pickAndModifyParams } from './model/queue/quickPick'
@@ -34,6 +34,8 @@ import { EventName } from '../telemetry/constants'
 import { Toast } from '../vscode/toast'
 import { getInput } from '../vscode/inputBox'
 import { createTypedAccumulator } from '../util/object'
+import { setContextValue } from '../vscode/context'
+import { standardizePath } from '../fileSystem/path'
 
 export const ExperimentsScale = {
   ...ColumnType,
@@ -42,6 +44,7 @@ export const ExperimentsScale = {
 } as const
 
 export class Experiments extends BaseRepository<TableData> {
+  public readonly onDidChangeIsParamsFileFocused: Event<string | undefined>
   public readonly onDidChangeExperiments: Event<ExperimentsOutput | void>
   public readonly onDidChangeColumns: Event<void>
   public readonly onDidChangeCheckpoints: Event<void>
@@ -55,6 +58,10 @@ export class Experiments extends BaseRepository<TableData> {
   private readonly columns: ColumnsModel
   private readonly checkpoints: CheckpointsModel
 
+  private readonly paramsFileFocused = this.dispose.track(
+    new EventEmitter<string | undefined>()
+  )
+
   private readonly experimentsChanged = this.dispose.track(
     new EventEmitter<ExperimentsOutput | void>()
   )
@@ -80,6 +87,7 @@ export class Experiments extends BaseRepository<TableData> {
 
     this.internalCommands = internalCommands
 
+    this.onDidChangeIsParamsFileFocused = this.paramsFileFocused.event
     this.onDidChangeExperiments = this.experimentsChanged.event
     this.onDidChangeColumns = this.columnsChanged.event
     this.onDidChangeCheckpoints = this.checkpointsChanged.event
@@ -113,6 +121,7 @@ export class Experiments extends BaseRepository<TableData> {
 
     this.handleMessageFromWebview()
     this.setupInitialData()
+    this.setActiveEditorContext()
   }
 
   public update() {
@@ -537,4 +546,44 @@ export class Experiments extends BaseRepository<TableData> {
 
     return experiment?.id
   }
+
+  private setActiveEditorContext() {
+    const setActiveEditorContext = (active: boolean) => {
+      setContextValue('dvc.params.fileActive', active)
+      const activeDvcRoot = active ? this.dvcRoot : undefined
+      this.paramsFileFocused.fire(activeDvcRoot)
+    }
+
+    this.dispose.track(
+      this.onDidChangeColumns(() => {
+        const path = standardizePath(window.activeTextEditor?.document.fileName)
+        if (!path) {
+          return
+        }
+
+        if (!this.columns.getParamsFiles().has(path)) {
+          return
+        }
+        setActiveEditorContext(true)
+      })
+    )
+
+    this.dispose.track(
+      window.onDidChangeActiveTextEditor(event => {
+        const path = standardizePath(event?.document.fileName)
+        if (!path) {
+          setActiveEditorContext(false)
+          return
+        }
+
+        if (path.includes(this.dvcRoot)) {
+          if (this.columns.getParamsFiles().has(path)) {
+            setActiveEditorContext(true)
+            return
+          }
+          setActiveEditorContext(false)
+        }
+      })
+    )
+  }
 }

extension/src/experiments/workspace.ts

@@ -24,6 +24,8 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
 
   private readonly checkpointsChanged: EventEmitter<void>
 
+  private focusedParamsDvcRoot: string | undefined
+
   constructor(
     internalCommands: InternalCommands,
     updatesPaused: EventEmitter<boolean>,
@@ -232,6 +234,13 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
         dvcRoot => (this.focusedWebviewDvcRoot = dvcRoot)
       )
     )
+
+    experiments.dispose.track(
+      experiments.onDidChangeIsParamsFileFocused(
+        dvcRoot => (this.focusedParamsDvcRoot = dvcRoot)
+      )
+    )
+
     experiments.dispose.track(
       experiments.onDidChangeExperiments(() => {
         this.experimentsChanged.fire()
@@ -253,6 +262,14 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
     return experiments
   }
 
+  public getFocusedOrOnlyOrPickProject() {
+    return (
+      this.focusedWebviewDvcRoot ||
+      this.focusedParamsDvcRoot ||
+      this.getOnlyOrPickProject()
+    )
+  }
+
   private async pickExpThenRun(
     commandId: CommandId,
     pickFunc: (

extension/src/fileSystem/path.ts

@@ -0,0 +1,8 @@
+import { Uri } from 'vscode'
+
+export const standardizePath = (path?: string): string | undefined => {
+  if (!path) {
+    return
+  }
+  return Uri.file(path).fsPath
+}

extension/src/persistence/model.ts

@@ -3,7 +3,7 @@ import { PersistenceKey } from './constants'
 import { DeferredDisposable } from '../class/deferred'
 
 export class ModelWithPersistence extends DeferredDisposable {
-  private readonly dvcRoot: string
+  protected readonly dvcRoot: string
   private readonly workspaceState: Memento
 
   constructor(dvcRoot: string, workspaceState: Memento) {

extension/src/plots/workspace.ts

@@ -46,4 +46,8 @@ export class WorkspacePlots extends BaseWorkspaceWebviews<Plots, PlotsData> {
     }
     return this.getRepository(dvcRoot).selectPlots()
   }
+
+  public getFocusedOrOnlyOrPickProject() {
+    return this.focusedWebviewDvcRoot || this.getOnlyOrPickProject()
+  }
 }