Add jupyter notebook script to run in dvc.yaml (#3240)

* Add jupyter notebook script to run in dvc.yaml

* Add tests

* Edit tests for windows

* Apply review comments

Author: sroy3

extension/src/fileSystem/index.test.ts

@@ -1,17 +1,31 @@
 import { join, relative, resolve } from 'path'
-import { ensureDirSync, remove } from 'fs-extra'
+import { appendFileSync, ensureDirSync, ensureFileSync, remove } from 'fs-extra'
 import {
   exists,
   findAbsoluteDvcRootPath,
   findDvcRootPaths,
   isDirectory,
   isSameOrChild,
-  getModifiedTime
+  getModifiedTime,
+  findOrCreateDvcYamlFile,
+  scriptCommand
 } from '.'
 import { dvcDemoPath } from '../test/util'
 import { DOT_DVC } from '../cli/dvc/constants'
 
 jest.mock('../cli/dvc/reader')
+jest.mock('fs-extra', () => {
+  const actualModule = jest.requireActual('fs-extra')
+  return {
+    __esModule: true,
+    ...actualModule,
+    appendFileSync: jest.fn(),
+    ensureFileSync: jest.fn()
+  }
+})
+
+const mockedAppendFileSync = jest.mocked(appendFileSync)
+const mockedEnsureFileSync = jest.mocked(ensureFileSync)
 
 beforeEach(() => {
   jest.resetAllMocks()
@@ -151,3 +165,70 @@ describe('getModifiedTime', () => {
     expect(epoch).toBeGreaterThan(1640995200000)
   })
 })
+
+describe('findOrCreateDvcYamlFile', () => {
+  it('should make sure a dvc.yaml file exists', () => {
+    const cwd = '/cwd'
+    findOrCreateDvcYamlFile(cwd, '/my/training/script.py')
+
+    expect(mockedEnsureFileSync).toHaveBeenCalledWith(`${cwd}/dvc.yaml`)
+  })
+
+  it('should add the training script as a train stage in the dvc.yaml file', () => {
+    const cwd = '/cwd'
+    findOrCreateDvcYamlFile(cwd, '/my/training/script.py')
+
+    expect(mockedAppendFileSync).toHaveBeenCalledWith(
+      `${cwd}/dvc.yaml`,
+      expect.stringMatching(/^\s+stages:\s+train:/)
+    )
+  })
+
+  it('should add the training script as a relative path to the cwd', () => {
+    findOrCreateDvcYamlFile(
+      '/dir/my_project/',
+      '/dir/my_project/src/training/train.py'
+    )
+
+    expect(mockedAppendFileSync).toHaveBeenCalledWith(
+      expect.anything(),
+      expect.stringContaining(join('src', 'training', 'train.py'))
+    )
+
+    findOrCreateDvcYamlFile(
+      '/dir/my_project/',
+      '/dir/my_other_project/train.py'
+    )
+
+    expect(mockedAppendFileSync).toHaveBeenCalledWith(
+      expect.anything(),
+      expect.stringContaining(join('..', 'my_other_project', 'train.py'))
+    )
+  })
+
+  it('should use the jupyter nbconvert command if the training script is a Jupyter notebook', () => {
+    findOrCreateDvcYamlFile('/', '/train.ipynb')
+
+    expect(mockedAppendFileSync).toHaveBeenCalledWith(
+      expect.anything(),
+      expect.stringContaining(scriptCommand.JUPYTER)
+    )
+    expect(mockedAppendFileSync).not.toHaveBeenCalledWith(
+      expect.anything(),
+      expect.stringContaining(scriptCommand.PYTHON)
+    )
+  })
+
+  it('should use the python command if the training script is not a Jupyter notebook', () => {
+    findOrCreateDvcYamlFile('/', '/train.py')
+
+    expect(mockedAppendFileSync).not.toHaveBeenCalledWith(
+      expect.anything(),
+      expect.stringContaining(scriptCommand.JUPYTER)
+    )
+    expect(mockedAppendFileSync).toHaveBeenCalledWith(
+      expect.anything(),
+      expect.stringContaining(scriptCommand.PYTHON)
+    )
+  })
+})

extension/src/fileSystem/index.ts

@@ -1,4 +1,4 @@
-import { basename, extname, join, relative, resolve, sep } from 'path'
+import { basename, extname, join, parse, relative, resolve, sep } from 'path'
 import {
   appendFileSync,
   ensureFileSync,
@@ -131,16 +131,25 @@ export const isAnyDvcYaml = (path?: string): boolean =>
       basename(path) === 'dvc.yaml')
   )
 
+export const scriptCommand = {
+  JUPYTER: 'jupyter nbconvert --to notebook --inplace --execute',
+  PYTHON: 'python'
+}
+
 export const findOrCreateDvcYamlFile = (
   cwd: string,
   trainingScript: string
 ) => {
   const dvcYamlPath = `${cwd}/dvc.yaml`
   ensureFileSync(dvcYamlPath)
 
-  const pipeline = `stages:
+  const isNotebook = parse(trainingScript).ext === '.ipynb'
+  const command = isNotebook ? scriptCommand.JUPYTER : scriptCommand.PYTHON
+
+  const pipeline = `
+stages:
   train:
-    cmd: python ${relative(cwd, trainingScript)}`
+    cmd: ${command} ${relative(cwd, trainingScript)}`
   return appendFileSync(dvcYamlPath, pipeline)
 }