Use diff in repository class

extension/src/Repository/DecorationProvider.ts

@@ -14,14 +14,20 @@ import { isStringInEnum } from '../util'
 export type DecorationState = Record<Status, Set<string>>
 
 enum Status {
+  ADDED = 'added',
   DELETED = 'deleted',
   MODIFIED = 'modified',
-  NEW = 'new',
   NOT_IN_CACHE = 'notInCache',
   TRACKED = 'tracked'
 }
 
 export class DecorationProvider implements FileDecorationProvider {
+  private static DecorationAdded: FileDecoration = {
+    badge: 'A',
+    color: new ThemeColor('gitDecoration.addedResourceForeground'),
+    tooltip: 'DVC added'
+  }
+
   private static DecorationDeleted: FileDecoration = {
     badge: 'D',
     color: new ThemeColor('gitDecoration.deletedResourceForeground'),
@@ -34,12 +40,6 @@ export class DecorationProvider implements FileDecorationProvider {
     tooltip: 'DVC modified'
   }
 
-  private static DecorationNew: FileDecoration = {
-    badge: 'A',
-    color: new ThemeColor('gitDecoration.addedResourceForeground'),
-    tooltip: 'DVC added'
-  }
-
   private static DecorationNotInCache: FileDecoration = {
     badge: 'NC',
     color: new ThemeColor('gitDecoration.renamedResourceForeground'),
@@ -102,8 +102,8 @@ export class DecorationProvider implements FileDecorationProvider {
     if (this.state.deleted?.has(uri.fsPath)) {
       return DecorationProvider.DecorationDeleted
     }
-    if (this.state.new?.has(uri.fsPath)) {
-      return DecorationProvider.DecorationNew
+    if (this.state.added?.has(uri.fsPath)) {
+      return DecorationProvider.DecorationAdded
     }
     if (this.state.notInCache?.has(uri.fsPath)) {
       return DecorationProvider.DecorationNotInCache

extension/src/Repository/index.test.ts

@@ -4,8 +4,13 @@ import { Config } from '../Config'
 import { SourceControlManagement } from './views/SourceControlManagement'
 import { mocked } from 'ts-jest/utils'
 import { DecorationProvider } from './DecorationProvider'
-import { Repository, RepositoryState, Status } from '.'
-import { listDvcOnlyRecursive, ListOutput, status } from '../cli/reader'
+import { Repository, RepositoryState } from '.'
+import {
+  diff,
+  DiffOutput,
+  listDvcOnlyRecursive,
+  ListOutput
+} from '../cli/reader'
 import { getAllUntracked } from '../git'
 
 jest.mock('@hediet/std/disposable')
@@ -15,8 +20,8 @@ jest.mock('../cli/reader')
 jest.mock('../git')
 jest.mock('../fileSystem')
 
+const mockedDiff = mocked(diff)
 const mockedListDvcOnlyRecursive = mocked(listDvcOnlyRecursive)
-const mockedStatus = mocked(status)
 const mockedGetAllUntracked = mocked(getAllUntracked)
 
 const mockedSourceControlManagement = mocked(SourceControlManagement)
@@ -52,6 +57,14 @@ beforeEach(() => {
 describe('Repository', () => {
   const dvcRoot = resolve(__dirname, '..', '..', 'demo')
 
+  const emptyDiff = {
+    added: [],
+    modified: [],
+    deleted: [],
+    renamed: [],
+    'not in cache': []
+  }
+
   describe('ready', () => {
     it('should wait for the state to be ready before resolving', async () => {
       const logDir = 'logs'
@@ -67,16 +80,19 @@ describe('Repository', () => {
         { path: rawDataDir }
       ] as ListOutput[])
 
-      mockedStatus.mockResolvedValueOnce({
-        train: [
-          { 'changed deps': { 'data/MNIST': 'modified' } },
-          { 'changed outs': { 'model.pt': 'modified', logs: 'modified' } },
-          'always changed'
+      mockedDiff.mockResolvedValueOnce({
+        added: [],
+        deleted: [],
+        modified: [
+          { path: 'model.pt' },
+          { path: logDir },
+          { path: logAcc },
+          { path: logLoss },
+          { path: MNISTDataDir }
         ],
-        'data/MNIST/raw.dvc': [
-          { 'changed outs': { 'data/MNIST/raw': 'modified' } }
-        ]
-      } as Record<string, (Record<string, Record<string, string>> | string)[]>)
+        'not in cache': [],
+        renamed: []
+      } as DiffOutput)
 
       const untracked = new Set([
         resolve(dvcRoot, 'some', 'untracked', 'python.py')
@@ -91,7 +107,13 @@ describe('Repository', () => {
       const repository = new Repository(dvcRoot, config, decorationProvider)
       await repository.isReady()
 
-      const modified = new Set([resolve(dvcRoot, rawDataDir)])
+      const modified = new Set([
+        resolve(dvcRoot, rawDataDir),
+        resolve(dvcRoot, logDir),
+        resolve(dvcRoot, logAcc),
+        resolve(dvcRoot, logLoss),
+        resolve(dvcRoot, model)
+      ])
       const tracked = new Set([
         resolve(dvcRoot, logAcc),
         resolve(dvcRoot, logLoss),
@@ -108,18 +130,18 @@ describe('Repository', () => {
         pythonBinPath: undefined
       }
 
-      expect(mockedStatus).toBeCalledWith(expectedExecutionOptions)
+      expect(mockedDiff).toBeCalledWith(expectedExecutionOptions)
       expect(mockedGetAllUntracked).toBeCalledWith(dvcRoot)
       expect(mockedListDvcOnlyRecursive).toBeCalledWith(
         expectedExecutionOptions
       )
 
       expect(repository.getState()).toEqual(
         expect.objectContaining({
+          added: emptySet,
           dispose: Disposable.fn(),
           deleted: emptySet,
           notInCache: emptySet,
-          new: emptySet,
           modified,
           tracked,
           untracked
@@ -131,7 +153,7 @@ describe('Repository', () => {
   describe('resetState', () => {
     it('will not exclude changed outs from stages that are always changed', async () => {
       mockedListDvcOnlyRecursive.mockResolvedValueOnce([])
-      mockedStatus.mockResolvedValueOnce({})
+      mockedDiff.mockResolvedValueOnce(emptyDiff)
       mockedGetAllUntracked.mockResolvedValueOnce(new Set())
 
       const config = ({
@@ -150,18 +172,12 @@ describe('Repository', () => {
       const logLoss = join(logDir, 'loss.tsv')
       const model = 'model.pt'
 
-      mockedStatus.mockResolvedValueOnce({
-        train: [
-          {
-            'changed deps': { 'data/MNIST': 'modified', 'train.py': 'modified' }
-          },
-          { 'changed outs': { 'model.pt': 'deleted' } },
-          'always changed'
-        ],
-        'data/MNIST/raw.dvc': [
-          { 'changed outs': { 'data/MNIST/raw': 'deleted' } }
-        ]
-      } as Record<string, (Record<string, Record<string, string>> | string)[]>)
+      mockedDiff.mockResolvedValueOnce(({
+        added: [],
+        deleted: [{ path: model }, { path: dataDir }],
+        modified: [],
+        'not in cache': []
+      } as unknown) as DiffOutput)
 
       const emptySet = new Set<string>()
 
@@ -197,15 +213,15 @@ describe('Repository', () => {
         pythonBinPath: undefined
       }
 
-      expect(mockedStatus).toBeCalledWith(expectedExecutionOptions)
+      expect(mockedDiff).toBeCalledWith(expectedExecutionOptions)
       expect(mockedGetAllUntracked).toBeCalledWith(dvcRoot)
       expect(mockedListDvcOnlyRecursive).toBeCalledWith(
         expectedExecutionOptions
       )
 
       expect(repository.getState()).toEqual({
+        added: emptySet,
         dispose: Disposable.fn(),
-        new: emptySet,
         modified: emptySet,
         notInCache: emptySet,
         deleted,
@@ -216,7 +232,7 @@ describe('Repository', () => {
 
     it("should update the classes state and call it's dependents", async () => {
       mockedListDvcOnlyRecursive.mockResolvedValueOnce([])
-      mockedStatus.mockResolvedValueOnce({})
+      mockedDiff.mockResolvedValueOnce(emptyDiff)
       mockedGetAllUntracked.mockResolvedValueOnce(new Set())
 
       const config = ({
@@ -231,39 +247,20 @@ describe('Repository', () => {
       const logAcc = join(logDir, 'acc.tsv')
       const logLoss = join(logDir, 'loss.tsv')
       const dataDir = 'data'
-      const model = 'model.pt'
+      const model = 'model.pkl'
       mockedListDvcOnlyRecursive.mockResolvedValueOnce([
         { path: logAcc },
         { path: logLoss },
         { path: model },
         { path: dataDir }
       ] as ListOutput[])
 
-      mockedStatus.mockResolvedValueOnce({
-        prepare: [
-          { 'changed deps': { 'data/data.xml': Status.NOT_IN_CACHE } },
-          { 'changed outs': { 'data/prepared': Status.NOT_IN_CACHE } }
-        ],
-        featurize: [
-          { 'changed deps': { 'data/prepared': Status.NOT_IN_CACHE } },
-          { 'changed outs': { 'data/features': 'modified' } }
-        ],
-        train: [
-          { 'changed deps': { 'data/features': 'modified' } },
-          { 'changed outs': { 'model.pkl': 'deleted' } }
-        ],
-        evaluate: [
-          {
-            'changed deps': {
-              'data/features': 'modified',
-              'model.pkl': 'deleted'
-            }
-          }
-        ],
-        'data/data.xml.dvc': [
-          { 'changed outs': { 'data/data.xml': Status.NOT_IN_CACHE } }
-        ]
-      } as Record<string, (Record<string, Record<string, string>> | string)[]>)
+      mockedDiff.mockResolvedValueOnce(({
+        added: [],
+        modified: [{ path: 'data/features' }],
+        deleted: [{ path: model }],
+        'not in cache': [{ path: 'data/data.xml' }, { path: 'data/prepared' }]
+      } as unknown) as DiffOutput)
 
       const untracked = new Set([
         resolve(dvcRoot, 'some', 'untracked', 'python.py'),
@@ -279,8 +276,8 @@ describe('Repository', () => {
       const deleted = new Set([join(dvcRoot, 'model.pkl')])
       const modified = new Set([join(dvcRoot, 'data/features')])
       const notInCache = new Set([
-        join(dvcRoot, 'data/data.xml'),
-        join(dvcRoot, 'data/prepared')
+        join(dvcRoot, 'data/prepared'),
+        join(dvcRoot, 'data/data.xml')
       ])
       const tracked = new Set([
         resolve(dvcRoot, logAcc),
@@ -296,15 +293,15 @@ describe('Repository', () => {
         pythonBinPath: undefined
       }
 
-      expect(mockedStatus).toBeCalledWith(expectedExecutionOptions)
+      expect(mockedDiff).toBeCalledWith(expectedExecutionOptions)
       expect(mockedGetAllUntracked).toBeCalledWith(dvcRoot)
       expect(mockedListDvcOnlyRecursive).toBeCalledWith(
         expectedExecutionOptions
       )
 
       expect(repository.getState()).toEqual({
+        added: new Set(),
         dispose: Disposable.fn(),
-        new: new Set(),
         modified,
         notInCache,
         deleted,

extension/src/Repository/index.ts

@@ -7,15 +7,15 @@ import {
 } from './views/SourceControlManagement'
 import { DecorationProvider, DecorationState } from './DecorationProvider'
 import { Deferred } from '@hediet/std/synchronization'
-import { status, listDvcOnlyRecursive } from '../cli/reader'
+import { diff, listDvcOnlyRecursive } from '../cli/reader'
 import { dirname, join } from 'path'
 import { observable, makeObservable } from 'mobx'
 import { getExecutionOptions } from '../cli/execution'
 
 export enum Status {
+  ADDED = 'added',
   DELETED = 'deleted',
   MODIFIED = 'modified',
-  NEW = 'new',
   NOT_IN_CACHE = 'not in cache'
 }
 
@@ -36,18 +36,18 @@ export class RepositoryState
   implements DecorationState, SourceControlManagementState {
   public dispose = Disposable.fn()
 
-  public tracked: Set<string>
+  public added: Set<string>
   public deleted: Set<string>
   public modified: Set<string>
-  public new: Set<string>
   public notInCache: Set<string>
+  public tracked: Set<string>
   public untracked: Set<string>
 
   constructor() {
     this.tracked = new Set<string>()
     this.deleted = new Set<string>()
     this.modified = new Set<string>()
-    this.new = new Set<string>()
+    this.added = new Set<string>()
     this.notInCache = new Set<string>()
     this.untracked = new Set<string>()
   }
@@ -144,20 +144,30 @@ export class Repository {
     return Object.values(filteredStatusOutput).reduce(statusReducer, {})
   }
 
-  private async getStatus(): Promise<Partial<Record<Status, Set<string>>>> {
-    const options = getExecutionOptions(this.config, this.dvcRoot)
-    const statusOutput = (await status(options)) as StatusOutput
+  private getStatus(): Partial<Record<Status, Set<string>>> {
+    const statusOutput = {} as StatusOutput
 
     return this.reduceToChangedOutsStatuses(statusOutput)
   }
 
   public async updateStatus() {
-    const status = await this.getStatus()
+    this.getStatus()
+
+    const options = getExecutionOptions(this.config, this.dvcRoot)
+    const diffOutput = await diff(options)
 
-    this.state.modified = status.modified || new Set<string>()
-    this.state.deleted = status.deleted || new Set<string>()
-    this.state.new = status.new || new Set<string>()
-    this.state.notInCache = status['not in cache'] || new Set<string>()
+    this.state.added = new Set<string>(
+      diffOutput.added.map(entry => join(this.dvcRoot, entry.path))
+    )
+    this.state.modified = new Set<string>(
+      diffOutput.modified.map(entry => join(this.dvcRoot, entry.path))
+    )
+    this.state.deleted = new Set<string>(
+      diffOutput.deleted.map(entry => join(this.dvcRoot, entry.path))
+    )
+    this.state.notInCache = new Set<string>(
+      diffOutput['not in cache'].map(entry => join(this.dvcRoot, entry.path))
+    )
   }
 
   public async updateUntracked() {

extension/src/Repository/views/SourceControlManagement.test.ts

@@ -39,24 +39,24 @@ describe('SourceControlManagement', () => {
       expect(sourceControlManagement.getState()).toEqual([])
 
       const updatedState = ({
+        added: new Set(['/some/new/path']),
         deleted: new Set(['/some/deleted/path', '/some/other/deleted/path']),
         dispose: () => undefined,
-        new: new Set(['/some/new/path']),
         tracked: new Set(['/some/excluded/tracked/path'])
       } as unknown) as SourceControlManagementState
 
       sourceControlManagement.setState(updatedState)
 
       expect(sourceControlManagement.getState()).toEqual([
+        { resourceUri: Uri.file('/some/new/path'), contextValue: 'added' },
         {
           resourceUri: Uri.file('/some/deleted/path'),
           contextValue: 'deleted'
         },
         {
           resourceUri: Uri.file('/some/other/deleted/path'),
           contextValue: 'deleted'
-        },
-        { resourceUri: Uri.file('/some/new/path'), contextValue: 'new' }
+        }
       ])
 
       sourceControlManagement.setState(initialState)