Call exp show with multiple branches

extension/src/cli/dvc/constants.ts

@@ -84,7 +84,8 @@ export enum ExperimentFlag {
   NO_FETCH = '--no-fetch',
   QUEUE = '--queue',
   RESET = '--reset',
-  NUM_COMMIT = '-n'
+  NUM_COMMIT = '-n',
+  REV = '--rev'
 }
 
 export enum GcPreserveFlag {

extension/src/cli/dvc/contract.ts

@@ -90,11 +90,13 @@ export type ExecutorState = {
 export type ExpWithError = {
   rev: string
   name?: string
+  branch: string
 } & DvcError
 
 type ExpWithData = {
   rev: string
   name?: string
+  branch: string
   data: ExpData
 }
 

extension/src/cli/dvc/reader.test.ts

@@ -97,6 +97,7 @@ describe('CliReader', () => {
       const cliOutput = await dvcReader.expShow(cwd)
       expect(cliOutput).toStrictEqual([
         {
+          branch: '',
           error: { msg: unexpectedStderr, type: 'Caught error' },
           rev: EXPERIMENT_WORKSPACE_ID
         }
@@ -110,7 +111,9 @@ describe('CliReader', () => {
       )
 
       const cliOutput = await dvcReader.expShow(cwd)
-      expect(cliOutput).toStrictEqual([{ rev: EXPERIMENT_WORKSPACE_ID }])
+      expect(cliOutput).toStrictEqual([
+        { branch: '', rev: EXPERIMENT_WORKSPACE_ID }
+      ])
     })
   })
 

extension/src/cli/dvc/reader.ts

@@ -74,7 +74,9 @@ export class DvcReader extends DvcCli {
       Flag.JSON
     )
     if (isDvcError(output) || isEmpty(output)) {
-      return [{ rev: EXPERIMENT_WORKSPACE_ID, ...(output as DvcError) }]
+      return [
+        { branch: '', rev: EXPERIMENT_WORKSPACE_ID, ...(output as DvcError) }
+      ]
     }
     return output
   }

extension/src/cli/git/constants.ts

@@ -25,12 +25,12 @@ export enum Command {
 
 export enum Flag {
   ALL = '--all',
+  COUNT = '--count',
   DIRECTORIES = '-d',
   DIRECTORY = '--directory',
   DOT = '.',
   EXCLUDE_STANDARD = '--exclude-standard',
   FORCE = '-f',
-  FULL_HISTORY = '--full-history',
   GET_URL = '--get-url',
   HARD = '--hard',
   NAME_ONLY = '--name-only',

extension/src/cli/git/reader.ts

@@ -9,6 +9,7 @@ import { isDirectory } from '../../fileSystem'
 export const autoRegisteredCommands = {
   GIT_GET_BRANCHES: 'getBranches',
   GIT_GET_COMMIT_MESSAGES: 'getCommitMessages',
+  GIT_GET_CURRENT_BRANCH: 'getCurrentBranch',
   GIT_GET_NUM_COMMITS: 'getNumCommits',
   GIT_GET_REMOTE_URL: 'getRemoteUrl',
   GIT_GET_REPOSITORY_ROOT: 'getGitRepositoryRoot',
@@ -80,18 +81,13 @@ export class GitReader extends GitCli {
     return new Set([...files, ...dirs])
   }
 
-  public async getNumCommits(cwd: string) {
-    const options = getOptions(
-      cwd,
-      Command.REV_LIST,
-      Flag.FULL_HISTORY,
-      Flag.ALL
-    )
+  public async getNumCommits(cwd: string, branch: string) {
+    const options = getOptions(cwd, Command.REV_LIST, Flag.COUNT, branch)
     try {
-      const revisions = await this.executeProcess(options)
-      return trimAndSplit(revisions).length
+      const nbCommits = await this.executeProcess(options)
+      return Number.parseInt(nbCommits)
     } catch {
-      return ''
+      return 0
     }
   }
 
@@ -105,6 +101,19 @@ export class GitReader extends GitCli {
     }
   }
 
+  public async getCurrentBranch(cwd: string): Promise<string> {
+    const options = getOptions(cwd, Command.BRANCH)
+    try {
+      const branches = await this.executeProcess(options)
+      return trimAndSplit(branches)[0]
+        .replace('* (', '')
+        .replace(')', '')
+        .replace('HEAD detached at ', '')
+    } catch {
+      return ''
+    }
+  }
+
   private async getUntrackedDirectories(cwd: string): Promise<string[]> {
     const options = getOptions(
       cwd,

extension/src/experiments/columns/collect/index.test.ts

@@ -733,6 +733,7 @@ describe('collectRelativeMetricsFiles', () => {
     expect(
       collectRelativeMetricsFiles([
         {
+          branch: 'current',
           rev: EXPERIMENT_WORKSPACE_ID,
           error: { msg: 'I broke', type: 'not important' }
         }

extension/src/experiments/data/collect.test.ts

@@ -16,6 +16,7 @@ describe('collectFiles', () => {
   it('should handle an error being returned', () => {
     const workspaceOnly = [
       {
+        branch: 'current',
         error: { msg: 'bad things are happening', type: 'today' },
         rev: EXPERIMENT_WORKSPACE_ID
       }

extension/src/experiments/data/index.ts

@@ -7,7 +7,7 @@ import {
 import { getRelativePattern } from '../../fileSystem/relativePattern'
 import { createFileSystemWatcher } from '../../fileSystem/watcher'
 import { AvailableCommands, InternalCommands } from '../../commands/internal'
-import { ExpShowOutput } from '../../cli/dvc/contract'
+import { EXPERIMENT_WORKSPACE_ID, ExpShowOutput } from '../../cli/dvc/contract'
 import { BaseData } from '../../data'
 import { DOT_DVC, ExperimentFlag } from '../../cli/dvc/constants'
 import { gitPath } from '../../cli/git/constants'
@@ -41,17 +41,10 @@ export class ExperimentsData extends BaseData<ExpShowOutput> {
   }
 
   public async update(): Promise<void> {
-    const flags = this.experiments.getIsBranchesView()
-      ? [ExperimentFlag.ALL_BRANCHES]
-      : [
-          ExperimentFlag.NUM_COMMIT,
-          this.experiments.getNbOfCommitsToShow().toString()
-        ]
-    const data = await this.internalCommands.executeCommand<ExpShowOutput>(
-      AvailableCommands.EXP_SHOW,
-      this.dvcRoot,
-      ...flags
-    )
+    const isBranchesView = this.experiments.getIsBranchesView()
+    const data: ExpShowOutput = isBranchesView
+      ? ((await this.expShow([ExperimentFlag.ALL_BRANCHES])) as ExpShowOutput)
+      : await this.updateCommitsView()
 
     this.collectFiles(data)
 
@@ -62,6 +55,70 @@ export class ExperimentsData extends BaseData<ExpShowOutput> {
     this.collectedFiles = collectFiles(data, this.collectedFiles)
   }
 
+  private async getBranchesToShowWithCurrent() {
+    const currentBranch = await this.internalCommands.executeCommand<string>(
+      AvailableCommands.GIT_GET_CURRENT_BRANCH,
+      this.dvcRoot
+    )
+
+    const branches = this.experiments.getBranchesToShow()
+
+    if (!branches.includes(currentBranch)) {
+      branches.push(currentBranch)
+      this.experiments.setBranchesToShow(branches)
+    }
+    return { branches, currentBranch }
+  }
+
+  private async expShow(flags: (ExperimentFlag | string)[], branch?: string) {
+    const data = await this.internalCommands.executeCommand<ExpShowOutput>(
+      AvailableCommands.EXP_SHOW,
+      this.dvcRoot,
+      ...flags
+    )
+    return data.map(exp => ({ ...exp, branch }))
+  }
+
+  private async updateCommitsView() {
+    const data: ExpShowOutput = []
+
+    const { branches, currentBranch } =
+      await this.getBranchesToShowWithCurrent()
+
+    const output = []
+    for (const branch of branches) {
+      const branchFlags = [
+        ExperimentFlag.REV,
+        branch,
+        ExperimentFlag.NUM_COMMIT,
+        this.experiments.getNbOfCommitsToShow(branch).toString()
+      ]
+      output.push(
+        new Promise<void>((resolve, reject) => {
+          this.expShow(branchFlags, branch)
+            .then(output => {
+              const newData = output as ExpShowOutput
+              if (branch !== currentBranch) {
+                const workspaceIndex = newData.findIndex(
+                  exp => exp.rev === EXPERIMENT_WORKSPACE_ID
+                )
+                newData.splice(workspaceIndex, 1)
+              }
+              data.push(...newData)
+              resolve()
+            })
+            .catch(() => {
+              reject(new Error('Output failed'))
+            })
+        })
+      )
+    }
+
+    await Promise.all(output)
+
+    return data
+  }
+
   private async updateAvailableBranchesToSelect() {
     const allBranches = await this.internalCommands.executeCommand<string[]>(
       AvailableCommands.GIT_GET_BRANCHES,

extension/src/experiments/index.ts

@@ -594,10 +594,11 @@ export class Experiments extends BaseRepository<TableData> {
         ),
       () => this.addStage(),
       (branchesSelected: string[]) => this.selectBranches(branchesSelected),
-      () =>
+      (branch: string) =>
         this.internalCommands.executeCommand<number>(
           AvailableCommands.GIT_GET_NUM_COMMITS,
-          this.dvcRoot
+          this.dvcRoot,
+          branch
         ),
       () => this.data.update()
     )

extension/src/experiments/model/collect.ts

@@ -176,16 +176,17 @@ const collectExpState = (
   expState: ExpState,
   commitData: { [sha: string]: CommitData }
 ): Experiment | undefined => {
-  const { rev, name } = expState
+  const { rev, name, branch } = expState
   const label =
     rev === EXPERIMENT_WORKSPACE_ID
       ? EXPERIMENT_WORKSPACE_ID
       : name || shortenForLabel(rev)
   const id = name || label
 
-  const experiment: Experiment = { id, label }
+  const experiment: Experiment = { id, label } as unknown as Experiment
 
   const baseline = transformExpState(experiment, expState)
+  baseline.branch = branch
 
   if (rev === EXPERIMENT_WORKSPACE_ID && !name) {
     acc.workspace = baseline
@@ -247,7 +248,7 @@ const collectExpRange = (
   }
   const expState = revs[0]
 
-  const { name, rev } = expState
+  const { name, rev, branch } = expState
 
   const label =
     rev === EXPERIMENT_WORKSPACE_ID
@@ -256,6 +257,7 @@ const collectExpRange = (
 
   const experiment = transformExpState(
     {
+      branch,
       id: name || label,
       label
     },
@@ -317,7 +319,11 @@ export const collectExperiments = (
     experimentsByCommit: new Map(),
     hasCheckpoints: hasCheckpoints(output),
     runningExperiments: [],
-    workspace: { id: EXPERIMENT_WORKSPACE_ID, label: EXPERIMENT_WORKSPACE_ID }
+    workspace: {
+      branch: '',
+      id: EXPERIMENT_WORKSPACE_ID,
+      label: EXPERIMENT_WORKSPACE_ID
+    }
   }
 
   const commitData = getCommitData(commitsOutput)

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

@@ -318,17 +318,21 @@ describe('ExperimentsModel', () => {
   it('should set the number of commits to show correctly', () => {
     const model = new ExperimentsModel('', buildMockMemento())
 
-    model.setNbfCommitsToShow(42)
+    model.setNbfCommitsToShow(42, 'special-branch')
 
-    expect(model.getNbOfCommitsToShow()).toBe(42)
+    expect(model.getNbOfCommitsToShow('special-branch')).toBe(42)
   })
 
   it('should persist the number of commits to show when changing it', () => {
     const memento = buildMockMemento()
     const model = new ExperimentsModel('', memento)
 
-    model.setNbfCommitsToShow(42)
+    model.setNbfCommitsToShow(42, 'main')
 
-    expect(memento.get(PersistenceKey.NUMBER_OF_COMMITS_TO_SHOW)).toBe(42)
+    expect(memento.get(PersistenceKey.NUMBER_OF_COMMITS_TO_SHOW)).toStrictEqual(
+      {
+        main: 42
+      }
+    )
   })
 })

extension/src/experiments/model/index.ts

@@ -60,7 +60,7 @@ export class ExperimentsModel extends ModelWithPersistence {
   private availableColors: Color[]
   private coloredStatus: ColoredStatus
   private starredExperiments: StarredExperiments
-  private numberOfCommitsToShow: number
+  private numberOfCommitsToShow: Record<string, number>
   private isBranchesView: boolean
   private branchesToShow: string[] = []
   private availableBranchesToShow: string[] = []
@@ -92,11 +92,15 @@ export class ExperimentsModel extends ModelWithPersistence {
       PersistenceKey.EXPERIMENTS_STARS,
       {}
     )
-    this.numberOfCommitsToShow = this.revive<number>(
+    this.numberOfCommitsToShow = this.revive<Record<string, number>>(
       PersistenceKey.NUMBER_OF_COMMITS_TO_SHOW,
-      DEFAULT_NUM_OF_COMMITS_TO_SHOW
+      {}
     )
 
+    if (typeof this.numberOfCommitsToShow === 'number') {
+      this.numberOfCommitsToShow = {}
+    }
+
     this.branchesToShow = this.revive<string[]>(
       PersistenceKey.EXPERIMENTS_BRANCHES,
       []
@@ -410,12 +414,16 @@ export class ExperimentsModel extends ModelWithPersistence {
     return this.finishedRunning
   }
 
-  public setNbfCommitsToShow(numberOfCommitsToShow: number) {
-    this.numberOfCommitsToShow = numberOfCommitsToShow
+  public setNbfCommitsToShow(numberOfCommitsToShow: number, branch: string) {
+    this.numberOfCommitsToShow[branch] = numberOfCommitsToShow
     this.persistNbOfCommitsToShow()
   }
 
-  public getNbOfCommitsToShow() {
+  public getNbOfCommitsToShow(branch: string) {
+    return this.numberOfCommitsToShow[branch] || DEFAULT_NUM_OF_COMMITS_TO_SHOW
+  }
+
+  public getAllNbOfCommitsToShow() {
     return this.numberOfCommitsToShow
   }