diff --git a/extension/package.json b/extension/package.json index 315b8fe7ca..5150a2b967 100644 --- a/extension/package.json +++ b/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", diff --git a/extension/src/pipeline/context.ts b/extension/src/pipeline/context.ts new file mode 100644 index 0000000000..69334b03d5 --- /dev/null +++ b/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 +): 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) + ) +} diff --git a/extension/src/pipeline/index.ts b/extension/src/pipeline/index.ts index 7b6486d43b..9cbe94a9f5 100644 --- a/extension/src/pipeline/index.ts +++ b/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 + public readonly onDidFocusProject: Event private updated: EventEmitter + private focusedPipeline: string | undefined + private readonly pipelineFileFocused: EventEmitter = + this.dispose.track(new EventEmitter()) + + private readonly onDidFocusPipelineFile: Event = + this.pipelineFileFocused.event + + private projectFocused: EventEmitter = this.dispose.track( + new EventEmitter() + ) + private readonly dvcRoot: string private readonly data: PipelineData private readonly model: PipelineModel @@ -51,7 +64,10 @@ export class Pipeline extends DeferredDisposable { this.updated = this.dispose.track(new EventEmitter()) this.onDidUpdate = this.updated.event + this.onDidFocusProject = this.projectFocused.event + void this.initialize() + this.watchActiveEditor() } public hasPipeline() { @@ -59,6 +75,11 @@ export class Pipeline extends DeferredDisposable { } 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) + }) + ) + } } diff --git a/extension/src/pipeline/workspace.ts b/extension/src/pipeline/workspace.ts index 0ee8ac1cde..e8f5a01e62 100644 --- a/extension/src/pipeline/workspace.ts +++ b/extension/src/pipeline/workspace.ts @@ -13,6 +13,8 @@ import { getOnDidChangeExtensions, isInstalled } from '../vscode/extensions' export class WorkspacePipeline extends BaseWorkspace { 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 { 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 diff --git a/extension/src/test/suite/experiments/index.test.ts b/extension/src/test/suite/experiments/index.test.ts index 196c19c049..90c775433d 100644 --- a/extension/src/test/suite/experiments/index.test.ts +++ b/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' + ) }) }) diff --git a/extension/src/test/suite/pipeline/index.test.ts b/extension/src/test/suite/pipeline/index.test.ts index 5d97364a24..42943b96b1 100644 --- a/extension/src/test/suite/pipeline/index.test.ts +++ b/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 + ) + }) }) diff --git a/extension/src/test/suite/util.ts b/extension/src/test/suite/util.ts index 8b58cdf786..dc81bdbd77 100644 --- a/extension/src/test/suite/util.ts +++ b/extension/src/test/suite/util.ts @@ -327,3 +327,14 @@ export const waitForEditorText = async (): Promise => { } return waitForEditorText() } + +export const getActiveEditorUpdatedEvent = (disposer: Disposer) => + new Promise(resolve => { + const listener = disposer.track( + window.onDidChangeActiveTextEditor(() => { + resolve(undefined) + disposer.untrack(listener) + listener.dispose() + }) + ) + }) diff --git a/extension/src/vscode/context.ts b/extension/src/vscode/context.ts index da47e232e8..3c3e0fd221 100644 --- a/extension/src/vscode/context.ts +++ b/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',