Open dvc.yaml automatically after adding a pipeline (#3242)

* Open dvc.yaml automatically after adding a pipeline

* Fix lint error

* Add a comment to direct the user towards the dvc.yaml stages documentation

* Cut comment in two and use function to open document in editor

Author: sroy3

extension/src/experiments/workspace.ts

@@ -439,7 +439,7 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
       if (!trainingScript) {
         return false
       }
-      findOrCreateDvcYamlFile(cwd, trainingScript, stageName)
+      void findOrCreateDvcYamlFile(cwd, trainingScript, stageName)
     }
     return true
   }

extension/src/fileSystem/index.test.ts

@@ -1,5 +1,6 @@
 import { join, relative, resolve } from 'path'
 import { appendFileSync, ensureDirSync, ensureFileSync, remove } from 'fs-extra'
+import { TextDocument, window, workspace } from 'vscode'
 import {
   exists,
   findAbsoluteDvcRootPath,
@@ -26,6 +27,13 @@ jest.mock('fs-extra', () => {
 
 const mockedAppendFileSync = jest.mocked(appendFileSync)
 const mockedEnsureFileSync = jest.mocked(ensureFileSync)
+const mockedWorkspace = jest.mocked(workspace)
+const mockedWindow = jest.mocked(window)
+const mockedOpenTextDocument = jest.fn()
+const mockedShowTextDocument = jest.fn()
+
+mockedWorkspace.openTextDocument = mockedOpenTextDocument
+mockedWindow.showTextDocument = mockedShowTextDocument
 
 beforeEach(() => {
   jest.resetAllMocks()
@@ -169,7 +177,7 @@ describe('getModifiedTime', () => {
 describe('findOrCreateDvcYamlFile', () => {
   it('should make sure a dvc.yaml file exists', () => {
     const cwd = '/cwd'
-    findOrCreateDvcYamlFile(cwd, '/my/training/script.py', 'train')
+    void findOrCreateDvcYamlFile(cwd, '/my/training/script.py', 'train')
 
     expect(mockedEnsureFileSync).toHaveBeenCalledWith(`${cwd}/dvc.yaml`)
   })
@@ -187,16 +195,29 @@ describe('findOrCreateDvcYamlFile', () => {
 
   it('should add the training script as a train stage in the dvc.yaml file', () => {
     const cwd = '/cwd'
-    findOrCreateDvcYamlFile(cwd, '/my/training/script.py', 'train')
+    void findOrCreateDvcYamlFile(cwd, '/my/training/script.py', 'train')
 
     expect(mockedAppendFileSync).toHaveBeenCalledWith(
       `${cwd}/dvc.yaml`,
-      expect.stringMatching(/^\s+stages:\s+train:/)
+      expect.stringMatching(/stages:\s+train:/)
+    )
+  })
+
+  it('should add a comment to direct the user towards the dvc.yaml stages documentation', () => {
+    const cwd = '/cwd'
+    void findOrCreateDvcYamlFile(cwd, '/my/training/script.py', 'train')
+
+    expect(mockedAppendFileSync).toHaveBeenCalledWith(
+      `${cwd}/dvc.yaml`,
+      expect.stringContaining(
+        `# Read about DVC pipeline configuration (https://dvc.org/doc/user-guide/project-structure/dvcyaml-files#stages)
+# to customize your stages even more`
+      )
     )
   })
 
   it('should add the training script as a relative path to the cwd', () => {
-    findOrCreateDvcYamlFile(
+    void findOrCreateDvcYamlFile(
       '/dir/my_project/',
       '/dir/my_project/src/training/train.py',
       'train'
@@ -207,7 +228,7 @@ describe('findOrCreateDvcYamlFile', () => {
       expect.stringContaining(join('src', 'training', 'train.py'))
     )
 
-    findOrCreateDvcYamlFile(
+    void findOrCreateDvcYamlFile(
       '/dir/my_project/',
       '/dir/my_other_project/train.py',
       'train'
@@ -220,7 +241,7 @@ describe('findOrCreateDvcYamlFile', () => {
   })
 
   it('should use the jupyter nbconvert command if the training script is a Jupyter notebook', () => {
-    findOrCreateDvcYamlFile('/', '/train.ipynb', 'train')
+    void findOrCreateDvcYamlFile('/', '/train.ipynb', 'train')
 
     expect(mockedAppendFileSync).toHaveBeenCalledWith(
       expect.anything(),
@@ -233,7 +254,7 @@ describe('findOrCreateDvcYamlFile', () => {
   })
 
   it('should use the python command if the training script is not a Jupyter notebook', () => {
-    findOrCreateDvcYamlFile('/', '/train.py', 'train')
+    void findOrCreateDvcYamlFile('/', '/train.py', 'train')
 
     expect(mockedAppendFileSync).not.toHaveBeenCalledWith(
       expect.anything(),
@@ -244,4 +265,18 @@ describe('findOrCreateDvcYamlFile', () => {
       expect.stringContaining(scriptCommand.PYTHON)
     )
   })
+
+  it('should open the dvc.yaml file in the editor', () => {
+    mockedOpenTextDocument.mockResolvedValue({} as TextDocument)
+
+    void findOrCreateDvcYamlFile('/', '/train.py', 'train')
+
+    expect(mockedOpenTextDocument).toHaveBeenCalledWith(
+      expect.objectContaining({
+        authority: 'dvc.yaml',
+        path: '/',
+        scheme: 'file'
+      })
+    )
+  })
 })

extension/src/fileSystem/index.ts

@@ -10,7 +10,7 @@ import {
   writeFileSync
 } from 'fs-extra'
 import { load } from 'js-yaml'
-import { Uri } from 'vscode'
+import { Uri, workspace, window } from 'vscode'
 import { standardizePath } from './path'
 import { definedAndNonEmpty } from '../util/array'
 import { Logger } from '../common/logger'
@@ -136,6 +136,12 @@ export const scriptCommand = {
   PYTHON: 'python'
 }
 
+export const openFileInEditor = async (filePath: string) => {
+  const document = await workspace.openTextDocument(Uri.file(filePath))
+  await window.showTextDocument(document)
+  return document
+}
+
 export const findOrCreateDvcYamlFile = (
   cwd: string,
   trainingScript: string,
@@ -148,9 +154,13 @@ export const findOrCreateDvcYamlFile = (
   const command = isNotebook ? scriptCommand.JUPYTER : scriptCommand.PYTHON
 
   const pipeline = `
+# Read about DVC pipeline configuration (https://dvc.org/doc/user-guide/project-structure/dvcyaml-files#stages)
+# to customize your stages even more
 stages:
   ${stageName}:
     cmd: ${command} ${relative(cwd, trainingScript)}`
+
+  void openFileInEditor(dvcYamlPath)
   return appendFileSync(dvcYamlPath, pipeline)
 }
 

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

@@ -73,6 +73,7 @@ import { ConfigKey } from '../../../vscode/config'
 import { EXPERIMENT_WORKSPACE_ID } from '../../../cli/dvc/contract'
 import * as Time from '../../../util/time'
 import { AvailableCommands } from '../../../commands/internal'
+import { openFileInEditor } from '../../../fileSystem'
 
 suite('Experiments Test Suite', () => {
   const disposable = Disposable.fn()
@@ -153,10 +154,7 @@ suite('Experiments Test Suite', () => {
       const { experiments } = buildExperiments(disposable)
 
       const windowSpy = spy(window, 'createWebviewPanel')
-      const uri = Uri.file(resolve(dvcDemoPath, 'train.py'))
-
-      const document = await workspace.openTextDocument(uri)
-      await window.showTextDocument(document)
+      const document = await openFileInEditor(resolve(dvcDemoPath, 'train.py'))
 
       expect(window.activeTextEditor?.document).to.deep.equal(document)