Focus pipeline when dvc.yaml file is open in the active editor (#4273)

extension/package.json

@@ -420,7 +420,8 @@
       {
         "title": "Show Pipeline DAG",
         "command": "dvc.showPipelineDAG",
-        "category": "DVC"
+        "category": "DVC",
+        "icon": "$(symbol-class)"
       },
       {
         "title": "Show Plots",
@@ -1094,6 +1095,11 @@
           "group": "navigation@0",
           "when": "dvc.experiments.webview.active && dvc.experiment.running && dvc.commands.available"
         },
+        {
+          "command": "dvc.showPipelineDAG",
+          "group": "navigation@0",
+          "when": "dvc.pipeline.file.active && dvc.commands.available && dvc.project.available"
+        },
         {
           "command": "dvc.runExperiment",
           "group": "navigation@1",

extension/src/pipeline/context.ts

@@ -0,0 +1,44 @@
+import { dirname } from 'path'
+import { EventEmitter, window } from 'vscode'
+import { Disposable, Disposer } from '@hediet/std/disposable'
+import { ContextKey, setContextValue } from '../vscode/context'
+import { standardizePossiblePath } from '../fileSystem/path'
+
+const setContextOnDidChangeActiveEditor = (
+  setActiveEditorContext: (path: string) => void,
+  dvcRoot: string
+): Disposable =>
+  window.onDidChangeActiveTextEditor(event => {
+    const path = standardizePossiblePath(event?.document.fileName)
+    if (!path) {
+      setActiveEditorContext('')
+      return
+    }
+
+    if (!path.includes(dvcRoot)) {
+      return
+    }
+
+    setActiveEditorContext(path)
+  })
+
+export const setContextForEditorTitleIcons = (
+  dvcRoot: string,
+  disposer: (() => void) & Disposer,
+  pipelineFileFocused: EventEmitter<string | undefined>
+): void => {
+  const setActiveEditorContext = (path: string) => {
+    const pipeline = path.endsWith('dvc.yaml') ? dirname(path) : undefined
+    void setContextValue(ContextKey.PIPELINE_FILE_ACTIVE, !!pipeline)
+    pipelineFileFocused.fire(pipeline)
+  }
+
+  const activePath = window.activeTextEditor?.document?.fileName
+  if (activePath?.startsWith(dvcRoot)) {
+    setActiveEditorContext(activePath)
+  }
+
+  disposer.track(
+    setContextOnDidChangeActiveEditor(setActiveEditorContext, dvcRoot)
+  )
+}

extension/src/pipeline/index.ts

@@ -1,6 +1,7 @@
 import { join } from 'path'
 import { Event, EventEmitter } from 'vscode'
 import { appendFileSync, writeFileSync } from 'fs-extra'
+import { setContextForEditorTitleIcons } from './context'
 import { PipelineData } from './data'
 import { PipelineModel } from './model'
 import { DeferredDisposable } from '../class/deferred'
@@ -30,9 +31,21 @@ const getScriptCommand = (script: string) => {
 
 export class Pipeline extends DeferredDisposable {
   public onDidUpdate: Event<void>
+  public readonly onDidFocusProject: Event<string | undefined>
 
   private updated: EventEmitter<void>
 
+  private focusedPipeline: string | undefined
+  private readonly pipelineFileFocused: EventEmitter<string | undefined> =
+    this.dispose.track(new EventEmitter())
+
+  private readonly onDidFocusPipelineFile: Event<string | undefined> =
+    this.pipelineFileFocused.event
+
+  private projectFocused: EventEmitter<string | undefined> = this.dispose.track(
+    new EventEmitter()
+  )
+
   private readonly dvcRoot: string
   private readonly data: PipelineData
   private readonly model: PipelineModel
@@ -51,14 +64,22 @@ export class Pipeline extends DeferredDisposable {
     this.updated = this.dispose.track(new EventEmitter<void>())
     this.onDidUpdate = this.updated.event
 
+    this.onDidFocusProject = this.projectFocused.event
+
     void this.initialize()
+    this.watchActiveEditor()
   }
 
   public hasPipeline() {
     return this.model.hasPipeline()
   }
 
   public async getCwd() {
+    const focusedPipeline = this.getFocusedPipeline()
+    if (focusedPipeline) {
+      return focusedPipeline
+    }
+
     await this.checkOrAddPipeline()
 
     const pipelines = this.model.getPipelines()
@@ -192,4 +213,23 @@ export class Pipeline extends DeferredDisposable {
   private writeDag(dag: string) {
     writeFileSync(join(this.dvcRoot, TEMP_DAG_FILE), dag)
   }
+
+  private getFocusedPipeline() {
+    return this.focusedPipeline
+  }
+
+  private watchActiveEditor() {
+    setContextForEditorTitleIcons(
+      this.dvcRoot,
+      this.dispose,
+      this.pipelineFileFocused
+    )
+
+    this.dispose.track(
+      this.onDidFocusPipelineFile(cwd => {
+        this.focusedPipeline = cwd
+        this.projectFocused.fire(cwd && this.dvcRoot)
+      })
+    )
+  }
 }

extension/src/pipeline/workspace.ts

@@ -13,6 +13,8 @@ import { getOnDidChangeExtensions, isInstalled } from '../vscode/extensions'
 export class WorkspacePipeline extends BaseWorkspace<Pipeline> {
   private isMermaidSupportInstalled = isInstalled(MARKDOWN_MERMAID_EXTENSION_ID)
 
+  private focusedProject: string | undefined
+
   constructor(internalCommands: InternalCommands) {
     super(internalCommands)
 
@@ -37,11 +39,17 @@ export class WorkspacePipeline extends BaseWorkspace<Pipeline> {
 
     this.setRepository(dvcRoot, pipeline)
 
+    this.dispose.track(
+      pipeline.onDidFocusProject(project => {
+        this.focusedProject = project
+      })
+    )
+
     return pipeline
   }
 
   public async showDag() {
-    const cwd = await this.getOnlyOrPickProject()
+    const cwd = this.focusedProject || (await this.getOnlyOrPickProject())
 
     if (!cwd) {
       return

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

@@ -47,6 +47,7 @@ import {
   configurationChangeEvent,
   experimentsUpdatedEvent,
   extensionUri,
+  getActiveEditorUpdatedEvent,
   getInputBoxEvent,
   getMessageReceivedEmitter
 } from '../util'
@@ -1893,17 +1894,6 @@ suite('Experiments Test Suite', () => {
   })
 
   describe('editor/title icons', () => {
-    const getActiveEditorUpdatedEvent = () =>
-      new Promise(resolve => {
-        const listener = disposable.track(
-          window.onDidChangeActiveTextEditor(() => {
-            resolve(undefined)
-            disposable.untrack(listener)
-            listener.dispose()
-          })
-        )
-      })
-
     it('should set the appropriate context value when a params file is open in the active editor/closed', async () => {
       const paramsFile = Uri.file(join(dvcDemoPath, 'params.yaml'))
       await window.showTextDocument(paramsFile)
@@ -1928,7 +1918,7 @@ suite('Experiments Test Suite', () => {
 
       mockSetContextValue.resetHistory()
 
-      const startupEditorClosed = getActiveEditorUpdatedEvent()
+      const startupEditorClosed = getActiveEditorUpdatedEvent(disposable)
 
       await closeAllEditors()
       await startupEditorClosed
@@ -1940,12 +1930,12 @@ suite('Experiments Test Suite', () => {
 
       mockSetContextValue.resetHistory()
 
-      const activeEditorUpdated = getActiveEditorUpdatedEvent()
+      const activeEditorUpdated = getActiveEditorUpdatedEvent(disposable)
 
       await window.showTextDocument(paramsFile)
       await activeEditorUpdated
 
-      const activeEditorClosed = getActiveEditorUpdatedEvent()
+      const activeEditorClosed = getActiveEditorUpdatedEvent(disposable)
 
       expect(
         mockContext['dvc.experiments.file.active'],
@@ -1970,7 +1960,9 @@ suite('Experiments Test Suite', () => {
       const { experiments } = buildExperiments({ disposer: disposable })
       await experiments.isReady()
 
-      expect(setContextValueSpy).not.to.be.called
+      expect(setContextValueSpy).not.to.be.calledWith(
+        'dvc.experiments.file.active'
+      )
     })
   })
 

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

@@ -1,12 +1,13 @@
 import { join } from 'path'
 import { afterEach, beforeEach, describe, it, suite } from 'mocha'
-import { SinonStub, restore, stub } from 'sinon'
+import { SinonStub, restore, spy, stub } from 'sinon'
 import { expect } from 'chai'
 import { QuickPickItem, Uri, window } from 'vscode'
 import { buildPipeline } from './util'
 import {
   bypassProcessManagerDebounce,
   closeAllEditors,
+  getActiveEditorUpdatedEvent,
   getMockNow
 } from '../util'
 import { Disposable } from '../../../extension'
@@ -15,6 +16,7 @@ import * as QuickPick from '../../../vscode/quickPick'
 import { QuickPickOptionsWithTitle } from '../../../vscode/quickPick'
 import * as FileSystem from '../../../fileSystem'
 import { ScriptCommand } from '../../../pipeline'
+import * as VscodeContext from '../../../vscode/context'
 
 suite('Pipeline Test Suite', () => {
   const disposable = Disposable.fn()
@@ -285,4 +287,76 @@ suite('Pipeline Test Suite', () => {
       expect(mockFindOrCreateDvcYamlFile).not.to.be.called
     })
   })
+
+  it('should set the appropriate context value when a dvc.yaml is open in the active editor', async () => {
+    const dvcYaml = Uri.file(join(dvcDemoPath, 'dvc.yaml'))
+    await window.showTextDocument(dvcYaml)
+
+    const mockContext: { [key: string]: unknown } = {
+      'dvc.pipeline.file.active': false
+    }
+
+    const mockSetContextValue = stub(VscodeContext, 'setContextValue')
+    mockSetContextValue.callsFake((key: string, value: unknown) => {
+      mockContext[key] = value
+      return Promise.resolve(undefined)
+    })
+
+    const { pipeline } = buildPipeline({ disposer: disposable })
+    await pipeline.isReady()
+
+    expect(
+      mockContext['dvc.pipeline.file.active'],
+      'should set dvc.pipeline.file.active to true when a dvc.yaml is open and the extension starts'
+    ).to.be.true
+
+    mockSetContextValue.resetHistory()
+
+    const startupEditorClosed = getActiveEditorUpdatedEvent(disposable)
+
+    await closeAllEditors()
+    await startupEditorClosed
+
+    expect(
+      mockContext['dvc.pipeline.file.active'],
+      'should set dvc.pipeline.file.active to false when the dvc.yaml in the active editor is closed'
+    ).to.be.false
+
+    mockSetContextValue.resetHistory()
+
+    const activeEditorUpdated = getActiveEditorUpdatedEvent(disposable)
+
+    await window.showTextDocument(dvcYaml)
+    await activeEditorUpdated
+
+    const activeEditorClosed = getActiveEditorUpdatedEvent(disposable)
+
+    expect(
+      mockContext['dvc.pipeline.file.active'],
+      'should set dvc.pipeline.file.active to true when a dvc.yaml file is in the active editor'
+    ).to.be.true
+
+    await closeAllEditors()
+    await activeEditorClosed
+
+    expect(
+      mockContext['dvc.pipeline.file.active'],
+      'should set dvc.pipeline.file.active to false when the dvc.yaml in the active editor is closed again'
+    ).to.be.false
+  })
+
+  it('should set dvc.pipeline.file.active to false when a dvc.yaml is not open and the extension starts', async () => {
+    const nonDvcYaml = Uri.file(join(dvcDemoPath, '.gitignore'))
+    await window.showTextDocument(nonDvcYaml)
+
+    const setContextValueSpy = spy(VscodeContext, 'setContextValue')
+
+    const { pipeline } = buildPipeline({ disposer: disposable })
+    await pipeline.isReady()
+
+    expect(setContextValueSpy).to.be.calledWith(
+      'dvc.pipeline.file.active',
+      false
+    )
+  })
 })

extension/src/test/suite/util.ts

@@ -327,3 +327,14 @@ export const waitForEditorText = async (): Promise<unknown> => {
   }
   return waitForEditorText()
 }
+
+export const getActiveEditorUpdatedEvent = (disposer: Disposer) =>
+  new Promise(resolve => {
+    const listener = disposer.track(
+      window.onDidChangeActiveTextEditor(() => {
+        resolve(undefined)
+        disposer.untrack(listener)
+        listener.dispose()
+      })
+    )
+  })

extension/src/vscode/context.ts

@@ -11,6 +11,7 @@ export enum ContextKey {
   EXPERIMENTS_SORTED = 'dvc.experiments.sorted',
   EXPERIMENTS_WEBVIEW_ACTIVE = 'dvc.experiments.webview.active',
   MULTIPLE_PROJECTS = 'dvc.multiple.projects',
+  PIPELINE_FILE_ACTIVE = 'dvc.pipeline.file.active',
   PLOTS_WEBVIEW_ACTIVE = 'dvc.plots.webview.active',
   PROJECT_AVAILABLE = 'dvc.project.available',
   PROJECT_HAS_DATA = 'dvc.project.hasData',