Fix plots file watchers (diff output key is not always a file)

extension/src/data/index.ts

@@ -108,7 +108,7 @@ export abstract class BaseData<
     )
   }
 
-  abstract collectFiles(data: T): string[]
-
   abstract managedUpdate(path?: string): Promise<unknown>
+
+  protected abstract collectFiles(data: T): string[]
 }

extension/src/experiments/data/index.ts

@@ -43,10 +43,6 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
     this.managedUpdate(QUEUED_EXPERIMENT_PATH)
   }
 
-  public collectFiles(data: ExperimentsOutput) {
-    return collectFiles(data, this.collectedFiles)
-  }
-
   public managedUpdate(path?: string) {
     if (
       path?.includes(QUEUED_EXPERIMENT_PATH) ||
@@ -76,6 +72,10 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
     return this.processManager.forceRunQueued()
   }
 
+  protected collectFiles(data: ExperimentsOutput) {
+    return collectFiles(data, this.collectedFiles)
+  }
+
   private async watchExpGitRefs(): Promise<void> {
     const gitRoot = await this.internalCommands.executeCommand(
       AvailableCommands.GIT_GET_REPOSITORY_ROOT,

extension/src/plots/data/collect.ts

@@ -0,0 +1,58 @@
+import { PlotsOutputOrError } from '../../cli/dvc/contract'
+import { isDvcError } from '../../cli/dvc/reader'
+import { uniqueValues } from '../../util/array'
+import { isImagePlot, Plot, TemplatePlot } from '../webview/contract'
+
+const collectImageFile = (acc: string[], file: string): void => {
+  if (acc.includes(file)) {
+    return
+  }
+  acc.push(file)
+}
+
+const collectFromDatapoint = (
+  acc: string[],
+  data: Record<string, unknown>
+): void => {
+  const filename = (data as { dvc_data_version_info?: { filename?: string } })
+    ?.dvc_data_version_info?.filename
+
+  if (!filename || acc.includes(filename)) {
+    return
+  }
+  acc.push(filename)
+}
+
+const collectTemplateFiles = (acc: string[], plot: TemplatePlot): void => {
+  for (const datapoints of Object.values(plot.datapoints || {})) {
+    for (const datapoint of datapoints) {
+      collectFromDatapoint(acc, datapoint)
+    }
+  }
+}
+
+const collectKeyData = (acc: string[], key: string, plots: Plot[]): void => {
+  for (const plot of plots) {
+    if (isImagePlot(plot)) {
+      collectImageFile(acc, key)
+      continue
+    }
+    collectTemplateFiles(acc, plot)
+  }
+}
+
+export const collectFiles = (
+  data: PlotsOutputOrError,
+  collectedFiles: string[]
+): string[] => {
+  const acc = [...collectedFiles]
+  if (isDvcError(data)) {
+    return acc
+  }
+
+  for (const [key, plots] of Object.entries(data)) {
+    collectKeyData(acc, key, plots)
+  }
+
+  return uniqueValues(acc)
+}

extension/src/plots/data/index.ts

@@ -1,6 +1,6 @@
 import { EventEmitter } from 'vscode'
+import { collectFiles } from './collect'
 import { PlotsOutputOrError } from '../../cli/dvc/contract'
-import { isDvcError } from '../../cli/dvc/reader'
 import { AvailableCommands, InternalCommands } from '../../commands/internal'
 import { BaseData } from '../../data'
 import {
@@ -59,22 +59,19 @@ export class PlotsData extends BaseData<{
       ...args
     )
 
+    this.notifyChanged({ data, revs })
+
     const files = this.collectFiles({ data })
 
     this.compareFiles(files)
-
-    return this.notifyChanged({ data, revs })
   }
 
   public managedUpdate() {
     return this.processManager.run('update')
   }
 
-  public collectFiles({ data }: { data: PlotsOutputOrError }) {
-    if (isDvcError(data)) {
-      return this.collectedFiles
-    }
-    return [...Object.keys(data), ...this.collectedFiles]
+  protected collectFiles({ data }: { data: PlotsOutputOrError }) {
+    return collectFiles(data, this.collectedFiles)
   }
 
   private getArgs(revs: string[]) {