Add file system watcher for dvc.yaml

extension/package.json

@@ -1063,6 +1063,7 @@
     "chokidar": "^3.5.2",
     "execa": "^5.1.1",
     "fs-extra": "^10.0.0",
+    "js-yaml": "^4.1.0",
     "lodash.clonedeep": "^4.5.0",
     "lodash.get": "^4.4.2",
     "lodash.isempty": "^4.4.0",
@@ -1078,6 +1079,7 @@
     "@types/copy-webpack-plugin": "^8.0.0",
     "@types/fs-extra": "^9.0.13",
     "@types/jest": "^27.4.0",
+    "@types/js-yaml": "^4.0.5",
     "@types/lodash.clonedeep": "^4.5.6",
     "@types/lodash.get": "^4.4.6",
     "@types/lodash.isempty": "^4.4.6",

extension/src/fileSystem/data/index.ts

@@ -0,0 +1,62 @@
+import { Event, EventEmitter } from 'vscode'
+import { Disposable } from '@hediet/std/disposable'
+import { Deferred } from '@hediet/std/synchronization'
+import { isSameOrChild, loadYaml, PartialDvcYaml } from '..'
+import { findFiles } from '../workspace'
+import { join } from '../../test/util/path'
+import { createFileSystemWatcher } from '../watcher'
+
+export class FileSystemData {
+  public readonly dispose = Disposable.fn()
+
+  public readonly onDidUpdate: Event<{ path: string; yaml: PartialDvcYaml }>
+
+  private readonly dvcRoot: string
+
+  private readonly updated = this.dispose.track(
+    new EventEmitter<{ path: string; yaml: PartialDvcYaml }>()
+  )
+
+  private readonly deferred = new Deferred()
+  private readonly initialized = this.deferred.promise
+
+  constructor(dvcRoot: string) {
+    this.dvcRoot = dvcRoot
+    this.onDidUpdate = this.updated.event
+
+    this.watchDvcYaml()
+    this.initialize()
+  }
+
+  public isReady() {
+    return this.initialized
+  }
+
+  private async initialize() {
+    const files = await findFiles(join('**', 'dvc.yaml'))
+    const filesInRepo = files.filter(file => isSameOrChild(this.dvcRoot, file))
+
+    filesInRepo.map(path => {
+      const yaml = loadYaml<PartialDvcYaml>(path)
+      if (yaml) {
+        this.updated.fire({ path, yaml })
+      }
+    })
+
+    this.deferred.resolve()
+  }
+
+  private watchDvcYaml() {
+    this.dispose.track(
+      createFileSystemWatcher(join(this.dvcRoot, '**', 'dvc.yaml'), path => {
+        if (!path) {
+          return
+        }
+        const yaml = loadYaml<PartialDvcYaml>(path)
+        if (yaml) {
+          this.updated.fire({ path, yaml })
+        }
+      })
+    )
+  }
+}

extension/src/fileSystem/index.ts

@@ -1,7 +1,15 @@
 import { basename, extname, join, relative, resolve } from 'path'
-import { existsSync, lstatSync, readdir, removeSync } from 'fs-extra'
+import {
+  existsSync,
+  lstatSync,
+  readdir,
+  readFileSync,
+  removeSync
+} from 'fs-extra'
+import { load } from 'js-yaml'
 import { Uri } from 'vscode'
 import { definedAndNonEmpty } from '../util/array'
+import { Logger } from '../common/logger'
 
 export const exists = (path: string): boolean => existsSync(path)
 
@@ -59,6 +67,7 @@ export type PartialDvcYaml = {
     train: { outs: (string | Record<string, { checkpoint?: boolean }>)[] }
   }
 }
+
 export const isAnyDvcYaml = (path?: string): boolean =>
   !!(
     path &&
@@ -71,3 +80,11 @@ export const relativeWithUri = (dvcRoot: string, uri: Uri) =>
   relative(dvcRoot, uri.fsPath)
 
 export const removeDir = (path: string): void => removeSync(path)
+
+export const loadYaml = <T>(path: string): T | undefined => {
+  try {
+    return load(readFileSync(path, 'utf-8')) as T
+  } catch {
+    Logger.error(`failed to load yaml ${path}`)
+  }
+}

extension/src/fileSystem/workspace.ts

@@ -25,3 +25,8 @@ export const isInWorkspace = (pathOrGlob: string): boolean => {
 
   return definedAndNonEmpty(isContained)
 }
+
+export const findFiles = async (relativeGlob: string): Promise<string[]> => {
+  const files = await workspace.findFiles(relativeGlob)
+  return files.map(uri => uri.fsPath)
+}

extension/src/test/suite/fileSystem/data/index.test.ts

@@ -0,0 +1,59 @@
+import { afterEach, beforeEach, describe, it, suite } from 'mocha'
+import { stub, restore } from 'sinon'
+import { expect } from 'chai'
+import { Disposable } from '../../../../extension'
+import { FileSystemData } from '../../../../fileSystem/data'
+import { dvcDemoPath } from '../../../util'
+import * as FileSystem from '../../../../fileSystem'
+import * as Watcher from '../../../../fileSystem/watcher'
+import { getFirstArgOfCall } from '../../util'
+import { join } from '../../../util/path'
+
+suite('Experiments Data Test Suite', () => {
+  const disposable = Disposable.fn()
+  const mockWatcher = {
+    dispose: stub()
+  } as Disposable
+
+  beforeEach(() => {
+    restore()
+  })
+
+  afterEach(() => {
+    disposable.dispose()
+  })
+
+  describe('FileSystemData', () => {
+    it('should read the dvc.yaml from the demo path and send an event containing the path and the yaml', async () => {
+      stub(Watcher, 'createFileSystemWatcher').returns(mockWatcher)
+      const data = disposable.track(new FileSystemData(dvcDemoPath))
+
+      disposable.track(
+        data.onDidUpdate(({ path, yaml }) => {
+          expect(path).to.equal(dvcDemoPath)
+          expect(yaml.stages.train.outs).to.deep.equal([
+            { 'model.pt': { checkpoint: true } }
+          ])
+        })
+      )
+
+      await data.isReady()
+    })
+
+    it('should create a watcher with the expected glob', async () => {
+      const mockCreateFileSystemWatcher = stub(
+        Watcher,
+        'createFileSystemWatcher'
+      ).returns(mockWatcher)
+      stub(FileSystem, 'loadYaml').returns(undefined)
+      const data = disposable.track(new FileSystemData(dvcDemoPath))
+
+      expect(mockCreateFileSystemWatcher).to.be.calledOnce
+      expect(getFirstArgOfCall(mockCreateFileSystemWatcher, 0)).to.equal(
+        join(dvcDemoPath, '**', 'dvc.yaml')
+      )
+
+      await data.isReady()
+    })
+  })
+})

yarn.lock

@@ -3458,6 +3458,11 @@
     jest-diff "^27.0.0"
     pretty-format "^27.0.0"
 
+"@types/js-yaml@^4.0.5":
+  version "4.0.5"
+  resolved "https://registry.yarnpkg.com/@types/js-yaml/-/js-yaml-4.0.5.tgz#738dd390a6ecc5442f35e7f03fa1431353f7e138"
+  integrity sha512-FhpRzf927MNQdRZP0J5DLIdTXhjLYzeUTmLAu69mnVksLH9CJY3IuSeEgbKUki7GQZm0WqDkGzyxju2EZGD2wA==
+
 "@types/jsdom@^16.2.6":
   version "16.2.13"
   resolved "https://registry.yarnpkg.com/@types/jsdom/-/jsdom-16.2.13.tgz#126c8b7441b159d6234610a48de77b6066f1823f"