Add dvc.yaml into files used to show editor/title icons

extension/package.json

@@ -1072,7 +1072,7 @@
  {
  "command": "dvc.stopAllRunningExperiments",
  "group": "navigation@0",
- "when": "dvc.params.file.active && dvc.experiment.running && dvc.commands.available"
+ "when": "dvc.experiments.file.active && dvc.experiment.running && dvc.commands.available"
  },
  {
  "command": "dvc.stopAllRunningExperiments",
@@ -1082,7 +1082,7 @@
  {
  "command": "dvc.runExperiment",
  "group": "navigation@1",
- "when": "dvc.params.file.active && !dvc.experiment.running.workspace && dvc.commands.available && !dvc.experiment.checkpoints"
+ "when": "dvc.experiments.file.active && !dvc.experiment.running.workspace && dvc.commands.available && !dvc.experiment.checkpoints"
  },
  {
  "command": "dvc.runExperiment",
@@ -1097,7 +1097,7 @@
  {
  "command": "dvc.resetAndRunCheckpointExperiment",
  "group": "navigation@2",
- "when": "dvc.params.file.active && !dvc.experiment.running.workspace && dvc.commands.available && dvc.experiment.checkpoints"
+ "when": "dvc.experiments.file.active && !dvc.experiment.running.workspace && dvc.commands.available && dvc.experiment.checkpoints"
  },
  {
  "command": "dvc.resumeCheckpointExperiment",
@@ -1107,7 +1107,7 @@
  {
  "command": "dvc.resumeCheckpointExperiment",
  "group": "navigation@3",
- "when": "dvc.params.file.active && !dvc.experiment.running.workspace && dvc.commands.available && dvc.experiment.checkpoints"
+ "when": "dvc.experiments.file.active && !dvc.experiment.running.workspace && dvc.commands.available && dvc.experiment.checkpoints"
  },
  {
  "command": "dvc.startExperimentsQueue",
@@ -1117,7 +1117,7 @@
  {
  "command": "dvc.startExperimentsQueue",
  "group": "navigation@4",
- "when": "dvc.params.file.active && dvc.commands.available"
+ "when": "dvc.experiments.file.active && dvc.commands.available"
  },
  {
  "command": "dvc.queueExperiment",
@@ -1127,7 +1127,7 @@
  {
  "command": "dvc.queueExperiment",
  "group": "navigation@5",
- "when": "dvc.params.file.active && !dvc.experiment.running.workspace && dvc.commands.available"
+ "when": "dvc.experiments.file.active && !dvc.experiment.running.workspace && dvc.commands.available"
  }
  ],
  "view/item/context": [

extension/src/experiments/context.ts

@@ -16,7 +16,7 @@ const setContextOnDidChangeParamsFiles = (
  return
  }
 
- if (!getParamsFiles().has(path)) {
+ if (!getParamsFiles().has(path) && !path.endsWith('dvc.yaml')) {
  return
  }
  setActiveEditorContext(true)
@@ -39,21 +39,25 @@ const setContextOnDidChangeActiveEditor = (
  }
 
  const isParamsFile = getParamsFiles().has(path)
+ const isDvcYaml = path.endsWith('dvc.yaml')
 
- setActiveEditorContext(isParamsFile)
+ setActiveEditorContext(isParamsFile || isDvcYaml)
  })
 
 export const setContextForEditorTitleIcons = (
  dvcRoot: string,
  disposer: (() => void) & Disposer,
  getParamsFiles: () => Set<string>,
- paramsFileFocused: EventEmitter<string | undefined>,
+ experimentsFileFocused: EventEmitter<string | undefined>,
  onDidChangeColumns: Event<void>
 ): void => {
- const setActiveEditorContext = (paramsFileActive: boolean) => {
- void setContextValue(ContextKey.PARAMS_FILE_ACTIVE, paramsFileActive)
- const activeDvcRoot = paramsFileActive ? dvcRoot : undefined
- paramsFileFocused.fire(activeDvcRoot)
+ const setActiveEditorContext = (experimentsFileActive: boolean) => {
+ void setContextValue(
+ ContextKey.EXPERIMENTS_FILE_ACTIVE,
+ experimentsFileActive
+ )
+ const activeDvcRoot = experimentsFileActive ? dvcRoot : undefined
+ experimentsFileFocused.fire(activeDvcRoot)
  }
 
  disposer.track(

extension/src/experiments/index.ts

@@ -56,7 +56,7 @@ export type ModifiedExperimentAndRunCommandId =
  | typeof AvailableCommands.EXPERIMENT_RESET_AND_RUN
 
 export class Experiments extends BaseRepository<TableData> {
- public readonly onDidChangeIsParamsFileFocused: Event<string | undefined>
+ public readonly onDidChangeIsExperimentsFileFocused: Event<string | undefined>
  public readonly onDidChangeExperiments: Event<void>
  public readonly onDidChangeColumns: Event<void>
  public readonly onDidChangeColumnOrderOrStatus: Event<void>
@@ -69,7 +69,7 @@ export class Experiments extends BaseRepository<TableData> {
  private readonly experiments: ExperimentsModel
  private readonly columns: ColumnsModel
 
- private readonly paramsFileFocused = this.dispose.track(
+ private readonly experimentsFileFocused = this.dispose.track(
  new EventEmitter<string | undefined>()
  )
 
@@ -123,7 +123,7 @@ export class Experiments extends BaseRepository<TableData> {
  this.addStage = addStage
  this.selectBranches = selectBranches
 
- this.onDidChangeIsParamsFileFocused = this.paramsFileFocused.event
+ this.onDidChangeIsExperimentsFileFocused = this.experimentsFileFocused.event
  this.onDidChangeExperiments = this.experimentsChanged.event
  this.onDidChangeColumns = this.columnsChanged.event
  this.onDidChangeColumnOrderOrStatus = this.columnsOrderOrStatusChanged.event
@@ -594,7 +594,7 @@ export class Experiments extends BaseRepository<TableData> {
  this.dvcRoot,
  this.dispose,
  () => this.columns.getParamsFiles(),
- this.paramsFileFocused,
+ this.experimentsFileFocused,
  this.onDidChangeColumns
  )
  }

extension/src/experiments/workspace.ts

@@ -65,7 +65,7 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
 
  private readonly checkpointsChanged: EventEmitter<void>
 
- private focusedParamsDvcRoot: string | undefined
+ private focusedFileDvcRoot: string | undefined
 
  constructor(
  internalCommands: InternalCommands,
@@ -327,8 +327,8 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
  )
 
  experiments.dispose.track(
- experiments.onDidChangeIsParamsFileFocused(
- dvcRoot => (this.focusedParamsDvcRoot = dvcRoot)
+ experiments.onDidChangeIsExperimentsFileFocused(
+ dvcRoot => (this.focusedFileDvcRoot = dvcRoot)
  )
  )
 
@@ -362,7 +362,7 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
  public getFocusedOrOnlyOrPickProject() {
  return (
  this.focusedWebviewDvcRoot ||
- this.focusedParamsDvcRoot ||
+ this.focusedFileDvcRoot ||
  this.getOnlyOrPickProject()
  )
  }

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

@@ -1962,7 +1962,7 @@ suite('Experiments Test Suite', () => {
  await window.showTextDocument(paramsFile)
 
  const mockContext: { [key: string]: unknown } = {
- 'dvc.params.file.active': false
+ 'dvc.experiments.file.active': false
  }
 
  const mockSetContextValue = stub(VscodeContext, 'setContextValue')
@@ -1975,8 +1975,8 @@ suite('Experiments Test Suite', () => {
  await experiments.isReady()
 
  expect(
- mockContext['dvc.params.file.active'],
- 'should set dvc.params.file.active to true when a params file is open and the extension starts'
+ mockContext['dvc.experiments.file.active'],
+ 'should set dvc.experiments.file.active to true when a params file is open and the extension starts'
  ).to.be.true
 
  mockSetContextValue.resetHistory()
@@ -1987,8 +1987,8 @@ suite('Experiments Test Suite', () => {
  await startupEditorClosed
 
  expect(
- mockContext['dvc.params.file.active'],
- 'should set dvc.params.file.active to false when the params file in the active editor is closed'
+ mockContext['dvc.experiments.file.active'],
+ 'should set dvc.experiments.file.active to false when the params file in the active editor is closed'
  ).to.be.false
 
  mockSetContextValue.resetHistory()
@@ -2001,16 +2001,16 @@ suite('Experiments Test Suite', () => {
  const activeEditorClosed = getActiveEditorUpdatedEvent()
 
  expect(
- mockContext['dvc.params.file.active'],
- 'should set dvc.params.file.active to true when a params file is in the active editor'
+ mockContext['dvc.experiments.file.active'],
+ 'should set dvc.experiments.file.active to true when a params file is in the active editor'
  ).to.be.true
 
  await closeAllEditors()
  await activeEditorClosed
 
  expect(
- mockContext['dvc.params.file.active'],
- 'should set dvc.params.file.active to false when the params file in the active editor is closed again'
+ mockContext['dvc.experiments.file.active'],
+ 'should set dvc.experiments.file.active to false when the params file in the active editor is closed again'
  ).to.be.false
  })
 

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

@@ -103,10 +103,15 @@ suite('Workspace Experiments Test Suite', () => {
  expect(mockQuickPickOne).to.not.be.called
  })
 
- it('should not prompt to pick a project if a params file is focused', async () => {
+ it('should not prompt to pick a project if a params file or dvc.yaml is focused', async () => {
  const mockQuickPickOne = stub(QuickPick, 'quickPickOne').resolves(
  dvcDemoPath
  )
+ const mockRunExperiment = stub(
+ DvcRunner.prototype,
+ 'runExperiment'
+ ).resolves(undefined)
+ stub(DvcReader.prototype, 'listStages').resolves('train')
 
  const { workspaceExperiments, experiments } =
  buildMultiRepoExperiments(disposable)
@@ -119,32 +124,37 @@ suite('Workspace Experiments Test Suite', () => {
 
  expect(await focusedWebview).to.equal(dvcDemoPath)
 
- const focusedParamsFile = new Promise(resolve => {
- const listener: Disposable = experiments.onDidChangeIsParamsFileFocused(
- (event: string | undefined) => {
- listener.dispose()
- return resolve(event)
- }
- )
- })
+ const getDvcRootFocusedEvent = () =>
+ new Promise(resolve => {
+ const listener: Disposable =
+ experiments.onDidChangeIsExperimentsFileFocused(
+ (event: string | undefined) => {
+ listener.dispose()
+ return resolve(event)
+ }
+ )
+ })
 
- const paramsFile = Uri.file(join(dvcDemoPath, 'params.yaml'))
- await window.showTextDocument(paramsFile)
+ const testFile = async (path: string) => {
+ const focusedDvcRoot = getDvcRootFocusedEvent()
+ const uri = Uri.file(join(dvcDemoPath, path))
+ await window.showTextDocument(uri)
 
- expect(await focusedParamsFile).to.equal(dvcDemoPath)
+  expect(await focusedDvcRoot).to.equal(dvcDemoPath)
 
- mockQuickPickOne.resetHistory()
+  mockQuickPickOne.resetHistory()
 
- const mockRunExperiment = stub(
- DvcRunner.prototype,
- 'runExperiment'
- ).resolves(undefined)
+ await workspaceExperiments.getCwdThenRun(
+ AvailableCommands.EXPERIMENT_RUN
+ )
 
- stub(DvcReader.prototype, 'listStages').resolves('train')
- await workspaceExperiments.getCwdThenRun(AvailableCommands.EXPERIMENT_RUN)
+ expect(mockQuickPickOne).not.to.be.called
+ expect(mockRunExperiment).to.be.calledWith(dvcDemoPath)
+ return closeAllEditors()
+ }
 
- expect(mockQuickPickOne).not.to.be.calledOnce
- expect(mockRunExperiment).to.be.calledWith(dvcDemoPath)
+ await testFile('params.yaml')
+ await testFile('dvc.yaml')
  })
  }).timeout(WEBVIEW_TEST_TIMEOUT)
 

extension/src/test/suite/vscode/recommend.test.ts

@@ -1,13 +1,22 @@
 import { join } from 'path'
-import { afterEach, beforeEach, describe, it, suite } from 'mocha'
+import { afterEach, before, beforeEach, describe, it, suite } from 'mocha'
 import { expect } from 'chai'
 import { MessageItem, Uri, window } from 'vscode'
 import { restore, stub } from 'sinon'
 import { closeAllEditors } from '../util'
 import { dvcDemoPath } from '../../util'
 import * as Extensions from '../../../vscode/extensions'
+import { recommendRedHatExtensionOnce } from '../../../vscode/recommend'
 
 suite('Recommend Test Suite', () => {
+ const openFileInEditor = (fileName: string) =>
+ window.showTextDocument(Uri.file(join(dvcDemoPath, fileName)))
+
+ before(async () => {
+ await openFileInEditor('dvc.lock') // clear any existing recommendation
+ return closeAllEditors()
+ })
+
  beforeEach(() => {
  restore()
  })
@@ -16,11 +25,9 @@ suite('Recommend Test Suite', () => {
  return closeAllEditors()
  })
 
- const openFileInEditor = (fileName: string) =>
- window.showTextDocument(Uri.file(join(dvcDemoPath, fileName)))
-
  describe('recommendRedHatExtensionOnce', () => {
  it('should only recommend the red hat yaml extension once per session', async () => {
+ recommendRedHatExtensionOnce()
  stub(Extensions, 'isInstalled').returns(false)
  const mockShowInformationMessage = stub(
  window,

extension/src/vscode/context.ts

@@ -6,11 +6,11 @@ export enum ContextKey {
  EXPERIMENT_CHECKPOINTS = 'dvc.experiment.checkpoints',
  EXPERIMENT_RUNNING = 'dvc.experiment.running',
  EXPERIMENT_RUNNING_WORKSPACE = 'dvc.experiment.running.workspace',
+ EXPERIMENTS_FILE_ACTIVE = 'dvc.experiments.file.active',
  EXPERIMENTS_FILTERED = 'dvc.experiments.filtered',
  EXPERIMENTS_SORTED = 'dvc.experiments.sorted',
  EXPERIMENTS_WEBVIEW_ACTIVE = 'dvc.experiments.webview.active',
  MULTIPLE_PROJECTS = 'dvc.multiple.projects',
- PARAMS_FILE_ACTIVE = 'dvc.params.file.active',
  PLOTS_WEBVIEW_ACTIVE = 'dvc.plots.webview.active',
  PROJECT_AVAILABLE = 'dvc.project.available',
  PROJECT_HAS_DATA = 'dvc.project.hasData',