Move experiments webview message sending into WebviewMessages (#2030)

* Move experiments webview message sending into WebviewMessages

* move webview messages creation out of constructor

* change get webview to a protected method

Author: mattseddon

extension/src/experiments/index.ts

@@ -73,6 +73,7 @@ export class Experiments extends BaseRepository<TableData> {
   )
 
   private readonly internalCommands: InternalCommands
+  private readonly webviewMessages: WebviewMessages
 
   constructor(
     dvcRoot: string,
@@ -119,7 +120,7 @@ export class Experiments extends BaseRepository<TableData> {
       })
     )
 
-    this.handleMessageFromWebview()
+    this.webviewMessages = this.createWebviewMessageHandler()
     this.setupInitialData()
     this.setActiveEditorContext()
   }
@@ -339,7 +340,7 @@ export class Experiments extends BaseRepository<TableData> {
   }
 
   public sendInitialWebviewData() {
-    return this.sendWebviewData()
+    return this.webviewMessages.sendWebviewMessage()
   }
 
   public getSelectedRevisions() {
@@ -433,43 +434,27 @@ export class Experiments extends BaseRepository<TableData> {
 
   private notifyColumnsChanged() {
     this.columnsChanged.fire()
-    return this.sendWebviewData()
-  }
-
-  private sendWebviewData() {
-    this.webview?.show(this.getWebviewData())
-  }
-
-  private getWebviewData() {
-    return {
-      changes: this.columns.getChanges(),
-      columnOrder: this.columns.getColumnOrder(),
-      columnWidths: this.columns.getColumnWidths(),
-      columns: this.columns.getSelected(),
-      filteredCounts: this.getFilteredCounts(),
-      filters: this.experiments.getFilterPaths(),
-      hasCheckpoints: this.hasCheckpoints(),
-      hasColumns: this.columns.hasColumns(),
-      hasRunningExperiment: this.experiments.hasRunningExperiment(),
-      rows: this.experiments.getRowData(),
-      sorts: this.experiments.getSorts()
-    }
+    return this.webviewMessages.sendWebviewMessage()
   }
 
-  private handleMessageFromWebview() {
+  private createWebviewMessageHandler() {
     const webviewMessages = new WebviewMessages(
       this.dvcRoot,
       this.experiments,
       this.columns,
+      this.checkpoints,
+      () => this.getWebview(),
       () => this.notifyChanged(),
       () => this.selectColumns()
     )
 
     this.dispose.track(
       this.onDidReceivedWebviewMessage(message =>
-        webviewMessages.handleMessageFromWebview(message)
+        this.webviewMessages.handleMessageFromWebview(message)
       )
     )
+
+    return webviewMessages
   }
 
   private async checkAutoApplyFilters(...filterIdsToRemove: string[]) {

extension/src/experiments/webview/messages.ts

@@ -1,4 +1,5 @@
 import { commands, Uri, ViewColumn, window } from 'vscode'
+import { TableData } from './contract'
 import {
   RegisteredCliCommands,
   RegisteredCommands
@@ -7,6 +8,7 @@ import { Logger } from '../../common/logger'
 import { sendTelemetryEvent } from '../../telemetry'
 import { EventName } from '../../telemetry/constants'
 import { join } from '../../test/util/path'
+import { BaseWebview } from '../../webview'
 import {
   MessageFromWebview,
   MessageFromWebviewType
@@ -15,30 +17,42 @@ import { ColumnsModel } from '../columns/model'
 import { splitColumnPath } from '../columns/paths'
 import { ExperimentsModel } from '../model'
 import { SortDefinition } from '../model/sortBy'
+import { CheckpointsModel } from '../checkpoints/model'
 
 export class WebviewMessages {
   private readonly dvcRoot: string
 
   private readonly experiments: ExperimentsModel
   private readonly columns: ColumnsModel
+  private readonly checkpoints: CheckpointsModel
 
+  private readonly getWebview: () => BaseWebview<TableData> | undefined
   private readonly notifyChanged: () => void
   private readonly selectColumns: () => void
 
   constructor(
     dvcRoot: string,
     experiments: ExperimentsModel,
     columns: ColumnsModel,
+    checkpoints: CheckpointsModel,
+    getWebview: () => BaseWebview<TableData> | undefined,
     notifyChanged: () => void,
     selectColumns: () => void
   ) {
     this.dvcRoot = dvcRoot
     this.experiments = experiments
     this.columns = columns
+    this.checkpoints = checkpoints
+    this.getWebview = getWebview
     this.notifyChanged = notifyChanged
     this.selectColumns = selectColumns
   }
 
+  public sendWebviewMessage() {
+    const webview = this.getWebview()
+    webview?.show(this.getWebviewData())
+  }
+
   public handleMessageFromWebview(message: MessageFromWebview) {
     switch (message.type) {
       case MessageFromWebviewType.REORDER_COLUMNS:
@@ -104,6 +118,24 @@ export class WebviewMessages {
     }
   }
 
+  private getWebviewData() {
+    return {
+      changes: this.columns.getChanges(),
+      columnOrder: this.columns.getColumnOrder(),
+      columnWidths: this.columns.getColumnWidths(),
+      columns: this.columns.getSelected(),
+      filteredCounts: this.experiments.getFilteredCounts(
+        this.checkpoints.hasCheckpoints()
+      ),
+      filters: this.experiments.getFilterPaths(),
+      hasCheckpoints: this.checkpoints.hasCheckpoints(),
+      hasColumns: this.columns.hasColumns(),
+      hasRunningExperiment: this.experiments.hasRunningExperiment(),
+      rows: this.experiments.getRowData(),
+      sorts: this.experiments.getSorts()
+    }
+  }
+
   private setColumnOrder(order: string[]) {
     this.columns.setColumnOrder(order)
     sendTelemetryEvent(

extension/src/webview/repository.ts

@@ -80,6 +80,10 @@ export abstract class BaseRepository<
     )
   }
 
+  protected getWebview() {
+    return this.webview
+  }
+
   private resetWebview() {
     this.isWebviewFocusedChanged.fire(undefined)
     this.dispose.untrack(this.webview)