Add filtered counts to Filter By tree (#1866)

* add filtered counts to the filter by tree description

* refactor text collection

* add integration test

Author: mattseddon

extension/src/experiments/index.ts

@@ -257,6 +257,10 @@ export class Experiments extends BaseRepository<TableData> {
     return this.notifyChanged()
   }
 
+  public getFilteredCounts() {
+    return this.experiments.getFilteredCounts()
+  }
+
   public getExperimentCount() {
     return this.experiments.getExperimentCount()
   }

extension/src/experiments/model/filterBy/tree.ts

@@ -3,19 +3,20 @@ import {
   ThemeIcon,
   TreeDataProvider,
   TreeItem,
-  TreeItemCollapsibleState
+  TreeItemCollapsibleState,
+  TreeView
 } from 'vscode'
 import { getFilterId } from '.'
 import { WorkspaceExperiments } from '../../workspace'
 import { RegisteredCommands } from '../../../commands/external'
 import { InternalCommands } from '../../../commands/internal'
 import { sendViewOpenedTelemetryEvent } from '../../../telemetry'
 import { EventName } from '../../../telemetry/constants'
-import { definedAndNonEmpty } from '../../../util/array'
+import { definedAndNonEmpty, joinTruthyItems } from '../../../util/array'
 import { createTreeView, getRootItem } from '../../../vscode/tree'
 import { Disposable } from '../../../class/dispose'
 
-type FilterItem = {
+export type FilterItem = {
   description: string
   dvcRoot: string
   id: string
@@ -29,6 +30,7 @@ export class ExperimentsFilterByTree
   public readonly onDidChangeTreeData: Event<string | void>
 
   private readonly experiments: WorkspaceExperiments
+  private readonly view: TreeView<string | FilterItem>
   private viewed = false
 
   constructor(
@@ -39,7 +41,7 @@ export class ExperimentsFilterByTree
 
     this.onDidChangeTreeData = experiments.experimentsChanged.event
 
-    this.dispose.track(
+    this.view = this.dispose.track(
       createTreeView<FilterItem>('dvc.views.experimentsFilterByTree', this)
     )
 
@@ -54,6 +56,8 @@ export class ExperimentsFilterByTree
       RegisteredCommands.EXPERIMENT_FILTERS_REMOVE_ALL,
       resource => this.removeAllFilters(resource)
     )
+
+    this.updateDescriptionOnChange()
   }
 
   public getTreeItem(element: string | FilterItem): TreeItem {
@@ -139,4 +143,58 @@ export class ExperimentsFilterByTree
   private getDvcRoots() {
     return this.experiments.getDvcRoots()
   }
+
+  private updateDescriptionOnChange() {
+    this.dispose.track(
+      this.onDidChangeTreeData(() => {
+        this.view.description = this.getDescription()
+      })
+    )
+  }
+
+  private getDescription() {
+    const dvcRoots = this.experiments.getDvcRoots()
+    if (!definedAndNonEmpty(dvcRoots)) {
+      return
+    }
+
+    const filtered = { checkpoints: 0, experiments: 0 }
+
+    for (const dvcRoot of dvcRoots) {
+      const { experiments, checkpoints } = this.experiments
+        .getRepository(dvcRoot)
+        .getFilteredCounts()
+      filtered.checkpoints = filtered.checkpoints + checkpoints
+      filtered.experiments = filtered.experiments + experiments
+    }
+
+    const combinedText = joinTruthyItems(
+      [
+        this.getDescriptionText('Experiment', filtered.experiments),
+        this.getDescriptionText('Checkpoint', filtered.checkpoints)
+      ],
+      ', '
+    )
+
+    if (!combinedText) {
+      return
+    }
+
+    return `${combinedText} Filtered`
+  }
+
+  private getDescriptionText(
+    type: 'Experiment' | 'Checkpoint',
+    filteredCount: number
+  ) {
+    if (!filteredCount) {
+      return
+    }
+
+    const text = `${filteredCount} ${type}`
+    if (filteredCount === 1) {
+      return text
+    }
+    return text + 's'
+  }
 }

extension/src/experiments/model/index.ts

@@ -348,6 +348,29 @@ export class ExperimentsModel extends ModelWithPersistence {
     ])
   }
 
+  public getFilteredCounts() {
+    const experiments = this.flattenExperiments()
+    const totalExperiments = experiments.length
+
+    const remainingExperiments = filterExperiments(
+      this.getFilters(),
+      experiments
+    ).length
+
+    const checkpoints = this.flattenCheckpoints()
+    const totalCheckpoints = checkpoints.length
+
+    const remainingCheckpoints = filterExperiments(
+      this.getFilters(),
+      checkpoints
+    ).length
+
+    return {
+      checkpoints: totalCheckpoints - remainingCheckpoints,
+      experiments: totalExperiments - remainingExperiments
+    }
+  }
+
   private getCombinedList() {
     return [
       this.workspace,

extension/src/test/suite/experiments/model/filterBy/tree.test.ts

@@ -1,8 +1,8 @@
 import { afterEach, beforeEach, describe, it, suite } from 'mocha'
 import { expect } from 'chai'
 import { stub, spy, restore } from 'sinon'
-import { window, commands, MessageItem } from 'vscode'
-import { addFilterViaQuickInput } from './util'
+import { window, commands, MessageItem, EventEmitter, TreeView } from 'vscode'
+import { addFilterViaQuickInput, mockQuickInputFilter } from './util'
 import { Disposable } from '../../../../../extension'
 import columnsFixture from '../../../../fixtures/expShow/columns'
 import rowsFixture from '../../../../fixtures/expShow/rows'
@@ -12,7 +12,7 @@ import {
   getFilterId,
   Operator
 } from '../../../../../experiments/model/filterBy'
-import { dvcDemoPath } from '../../../../util'
+import { buildMockMemento, dvcDemoPath } from '../../../../util'
 import { experimentsUpdatedEvent } from '../../../util'
 import { buildMetricOrParamPath } from '../../../../../experiments/columns/paths'
 import { RegisteredCommands } from '../../../../../commands/external'
@@ -25,6 +25,10 @@ import { WEBVIEW_TEST_TIMEOUT } from '../../../timeouts'
 import { Response } from '../../../../../vscode/response'
 import { ExperimentsModel } from '../../../../../experiments/model'
 import { Title } from '../../../../../vscode/title'
+import {
+  ExperimentsFilterByTree,
+  FilterItem
+} from '../../../../../experiments/model/filterBy/tree'
 
 suite('Experiments Filter By Tree Test Suite', () => {
   const disposable = Disposable.fn()
@@ -345,5 +349,105 @@ suite('Experiments Filter By Tree Test Suite', () => {
         'should not call get repository in removeFilters without a root'
       ).not.to.be.called
     })
+
+    it('should update the description when a filter is added or removed', async () => {
+      const { experiments, internalCommands } = buildExperiments(disposable)
+      await experiments.isReady()
+
+      const workspaceExperiments = disposable.track(
+        new WorkspaceExperiments(
+          internalCommands,
+          disposable.track(new EventEmitter()),
+          buildMockMemento(),
+          { [dvcDemoPath]: experiments },
+          disposable.track(new EventEmitter())
+        )
+      )
+      disposable.track(
+        experiments.onDidChangeExperiments(() =>
+          workspaceExperiments.experimentsChanged.fire()
+        )
+      )
+
+      const mockTreeView = {
+        description: undefined,
+        dispose: stub()
+      } as unknown as TreeView<string | FilterItem>
+      stub(window, 'createTreeView').returns(mockTreeView)
+
+      stub(internalCommands, 'registerExternalCommand').returns(undefined)
+      disposable.track(
+        new ExperimentsFilterByTree(workspaceExperiments, internalCommands)
+      )
+      const getUpdateEvent = () =>
+        new Promise(resolve =>
+          disposable.track(
+            workspaceExperiments.onDidChangeExperiments(() =>
+              resolve(undefined)
+            )
+          )
+        )
+
+      const tableFilterAdded = getUpdateEvent()
+
+      const filter = {
+        operator: Operator.EQUAL,
+        path: buildMetricOrParamPath(
+          ColumnType.METRICS,
+          'summary.json',
+          'loss'
+        ),
+        value: '0'
+      }
+
+      mockQuickInputFilter(filter)
+      experiments.addFilter()
+      await tableFilterAdded
+
+      expect(experiments.getFilteredCounts()).to.deep.equal({
+        checkpoints: 9,
+        experiments: 3
+      })
+
+      expect(mockTreeView.description).to.equal(
+        '3 Experiments, 9 Checkpoints Filtered'
+      )
+
+      const tableFilterRemoved = getUpdateEvent()
+      experiments.removeFilter(getFilterId(filter))
+
+      await tableFilterRemoved
+
+      expect(experiments.getFilteredCounts()).to.deep.equal({
+        checkpoints: 0,
+        experiments: 0
+      })
+
+      expect(mockTreeView.description).to.be.undefined
+
+      stub(experiments, 'getFilteredCounts')
+        .onFirstCall()
+        .returns({
+          checkpoints: 1,
+          experiments: 0
+        })
+        .onSecondCall()
+        .returns({
+          checkpoints: 0,
+          experiments: 1
+        })
+
+      const allButOneCheckpointFilteredEvent = getUpdateEvent()
+      workspaceExperiments.experimentsChanged.fire()
+      await allButOneCheckpointFilteredEvent
+
+      expect(mockTreeView.description).to.equal('1 Checkpoint Filtered')
+
+      const allButOneExperimentFilteredEvent = getUpdateEvent()
+      workspaceExperiments.experimentsChanged.fire()
+      await allButOneExperimentFilteredEvent
+
+      expect(mockTreeView.description).to.equal('1 Experiment Filtered')
+    })
   })
 })

extension/src/test/suite/experiments/model/filterBy/util.ts

@@ -6,15 +6,11 @@ import { experimentsUpdatedEvent } from '../../../util'
 import { Experiments } from '../../../../../experiments'
 import { RegisteredCommands } from '../../../../../commands/external'
 
-export const addFilterViaQuickInput = (
-  experiments: Experiments,
+export const mockQuickInputFilter = (
   fixtureFilter: FilterDefinition,
   mockShowQuickPick = stub(window, 'showQuickPick'),
   mockShowInputBox = stub(window, 'showInputBox')
 ) => {
-  mockShowQuickPick.resetHistory()
-  mockShowInputBox.resetHistory()
-
   const column = columnsFixture.find(
     column => column.path === fixtureFilter.path
   )
@@ -25,6 +21,18 @@ export const addFilterViaQuickInput = (
     value: fixtureFilter.operator
   } as unknown as QuickPickItem)
   mockShowInputBox.onFirstCall().resolves(fixtureFilter.value as string)
+}
+
+export const addFilterViaQuickInput = (
+  experiments: Experiments,
+  fixtureFilter: FilterDefinition,
+  mockShowQuickPick = stub(window, 'showQuickPick'),
+  mockShowInputBox = stub(window, 'showInputBox')
+) => {
+  mockShowQuickPick.resetHistory()
+  mockShowInputBox.resetHistory()
+
+  mockQuickInputFilter(fixtureFilter, mockShowQuickPick, mockShowInputBox)
 
   const tableFilterAdded = experimentsUpdatedEvent(experiments)