Make plots tree items decoratable (#3504)

iterative · Mar 20, 2023 · 4c54059 · 4c54059
1 parent a2ec645
commit 4c54059
Show file tree

Hide file tree

Showing 10 changed files with 44 additions and 48 deletions.
diff --git a/extension/src/experiments/columns/tree.ts b/extension/src/experiments/columns/tree.ts
@@ -1,3 +1,4 @@
+import { TreeItem } from 'vscode'
 import {
   BasePathSelectionTree,
   PathSelectionItem
@@ -30,14 +31,21 @@ export class ExperimentsColumnsTree extends BasePathSelectionTree<WorkspaceExper
     )
   }
 
-  public getRepositoryChildren(dvcRoot: string, path: string) {
+  protected getBaseTreeItem({
+    label,
+    collapsibleState
+  }: PathSelectionItem & { label: string }) {
+    return new TreeItem(label, collapsibleState)
+  }
+
+  protected getRepositoryChildren(dvcRoot: string, path: string) {
     return this.workspace
       .getRepository(dvcRoot)
       .getChildColumns(path)
       .map(element => this.transformElement({ ...element, dvcRoot }))
   }
 
-  public getRepositoryStatuses(dvcRoot: string) {
+  protected getRepositoryStatuses(dvcRoot: string) {
     return this.workspace.getRepository(dvcRoot).getColumnsStatuses()
   }
 }
diff --git a/extension/src/path/selection/model.ts b/extension/src/path/selection/model.ts
@@ -144,7 +144,7 @@ export abstract class PathSelectionModel<
 
   abstract getChildren(
     ...args: unknown[]
-  ): (T & { descendantStatuses: Status[]; label: string; status: Status })[]
+  ): (T & { descendantStatuses: Status[]; status: Status })[]
 
   abstract getTerminalNodes(): (T & { selected: boolean })[]
 }
diff --git a/extension/src/path/selection/tree.ts b/extension/src/path/selection/tree.ts
@@ -21,7 +21,7 @@ export type PathSelectionItem = {
   description: string | undefined
   dvcRoot: string
   collapsibleState: TreeItemCollapsibleState
-  label: string
+  label: string | undefined
   path: string
   iconPath: Resource | Uri
 }
@@ -76,10 +76,9 @@ export abstract class BasePathSelectionTree<
       return new TreeItem(resourceUri, TreeItemCollapsibleState.Collapsed)
     }
 
-    const { dvcRoot, path, label, collapsibleState, description, iconPath } =
-      element
+    const { dvcRoot, path, description, iconPath } = element
 
-    const treeItem = new TreeItem(label, collapsibleState)
+    const treeItem = this.getBaseTreeItem(element)
 
     treeItem.command = {
       arguments: [{ dvcRoot, path }],
@@ -132,7 +131,7 @@ export abstract class BasePathSelectionTree<
     hasChildren: boolean
     path: string
     status: Status
-    label: string
+    label?: string
   }) {
     const { dvcRoot, descendantStatuses, hasChildren, path, status, label } =
       element
@@ -203,10 +202,12 @@ export abstract class BasePathSelectionTree<
     return typeof element === 'string'
   }
 
-  abstract getRepositoryStatuses(dvcRoot: string): Status[]
+  protected abstract getBaseTreeItem(element: PathSelectionItem): TreeItem
 
-  abstract getRepositoryChildren(
+  protected abstract getRepositoryChildren(
     dvcRoot: string,
     path: string | undefined
   ): PathSelectionItem[]
+
+  protected abstract getRepositoryStatuses(dvcRoot: string): Status[]
 }
diff --git a/extension/src/plots/paths/collect.test.ts b/extension/src/plots/paths/collect.test.ts
@@ -25,61 +25,53 @@ describe('collectPath', () => {
     expect(collectPaths([], plotsDiffFixture, revisions, {})).toStrictEqual([
       {
         hasChildren: false,
-        label: 'acc.png',
         parentPath: 'plots',
         path: join('plots', 'acc.png'),
         revisions: new Set(revisions),
         type: new Set(['comparison'])
       },
       {
         hasChildren: true,
-        label: 'plots',
         parentPath: undefined,
         path: 'plots',
         revisions: new Set(revisions)
       },
       {
         hasChildren: false,
-        label: 'heatmap.png',
         parentPath: 'plots',
         path: join('plots', 'heatmap.png'),
         revisions: new Set(revisions),
         type: new Set(['comparison'])
       },
       {
         hasChildren: false,
-        label: 'loss.png',
         parentPath: 'plots',
         path: join('plots', 'loss.png'),
         revisions: new Set(revisions),
         type: new Set(['comparison'])
       },
       {
         hasChildren: false,
-        label: 'loss.tsv',
         parentPath: 'logs',
         path: join('logs', 'loss.tsv'),
         revisions: new Set(revisions),
         type: new Set(['template-single'])
       },
       {
         hasChildren: true,
-        label: 'logs',
         parentPath: undefined,
         path: 'logs',
         revisions: new Set(revisions)
       },
       {
         hasChildren: false,
-        label: 'acc.tsv',
         parentPath: 'logs',
         path: join('logs', 'acc.tsv'),
         revisions: new Set(revisions),
         type: new Set(['template-single'])
       },
       {
         hasChildren: false,
-        label: 'predictions.json',
         parentPath: undefined,
         path: 'predictions.json',
         revisions: new Set(revisions),
@@ -127,7 +119,6 @@ describe('collectPath', () => {
     const mockPath = 'completely:madeup:path'
     const mockPlotPath = {
       hasChildren: false,
-      label: mockPath,
       parentPath: undefined,
       path: mockPath,
       revisions: new Set(['bfc7f64']),
@@ -185,52 +176,45 @@ describe('collectPath', () => {
     expect(collectPaths([], mockPlotsDiff, revisions, {})).toStrictEqual([
       {
         hasChildren: false,
-        label: 'acc.tsv',
         parentPath: join('logs', 'scalars'),
         path: join('logs', 'scalars', 'acc.tsv'),
         revisions: new Set(revisions),
         type: new Set(['template-single'])
       },
       {
         hasChildren: true,
-        label: 'scalars',
         parentPath: 'logs',
         path: join('logs', 'scalars'),
         revisions: new Set(revisions)
       },
       {
         hasChildren: true,
-        label: 'logs',
         parentPath: undefined,
         path: 'logs',
         revisions: new Set(revisions)
       },
       {
         hasChildren: false,
-        label: 'loss.tsv',
         parentPath: join('logs', 'scalars'),
         path: join('logs', 'scalars', 'loss.tsv'),
         revisions: new Set(revisions),
         type: new Set(['template-single'])
       },
       {
         hasChildren: false,
-        label: 'heatmap.png',
         parentPath: 'plots',
         path: join('plots', 'heatmap.png'),
         revisions: new Set(revisions),
         type: new Set(['comparison'])
       },
       {
         hasChildren: true,
-        label: 'plots',
         parentPath: undefined,
         path: 'plots',
         revisions: new Set(revisions)
       },
       {
         hasChildren: false,
-        label: 'predictions.json',
         parentPath: undefined,
         path: 'predictions.json',
         revisions: new Set(revisions),

diff --git a/extension/src/plots/paths/collect.ts b/extension/src/plots/paths/collect.ts
@@ -25,7 +25,6 @@ export enum PathType {
 }
 
 export type PlotPath = {
-  label: string
   path: string
   type?: Set<PathType>
   parentPath: string | undefined
@@ -148,7 +147,6 @@ const collectOrderedPath = (
 
   const plotPath: PlotPath = {
     hasChildren,
-    label: pathArray[idx - 1],
     parentPath: getParent(pathArray, idx),
     path,
     revisions

diff --git a/extension/src/plots/paths/model.test.ts b/extension/src/plots/paths/model.test.ts
@@ -31,7 +31,6 @@ describe('PathsModel', () => {
     expect(model.getTerminalNodes()).toStrictEqual([
       {
         hasChildren: false,
-        label: 'acc.png',
         parentPath: 'plots',
         path: plotsAcc,
         revisions: new Set(revisions),
@@ -40,7 +39,6 @@ describe('PathsModel', () => {
       },
       {
         hasChildren: false,
-        label: 'heatmap.png',
         parentPath: 'plots',
         path: join('plots', 'heatmap.png'),
         revisions: new Set(revisions),
@@ -49,7 +47,6 @@ describe('PathsModel', () => {
       },
       {
         hasChildren: false,
-        label: 'loss.png',
         parentPath: 'plots',
         path: join('plots', 'loss.png'),
         revisions: new Set(revisions),
@@ -58,7 +55,6 @@ describe('PathsModel', () => {
       },
       {
         hasChildren: false,
-        label: 'loss.tsv',
         parentPath: 'logs',
         path: logsLoss,
         revisions: new Set(revisions),
@@ -67,7 +63,6 @@ describe('PathsModel', () => {
       },
       {
         hasChildren: false,
-        label: 'acc.tsv',
         parentPath: 'logs',
         path: logsAcc,
         revisions: new Set(revisions),
@@ -76,7 +71,6 @@ describe('PathsModel', () => {
       },
       {
         hasChildren: false,
-        label: 'predictions.json',
         parentPath: undefined,
         path: 'predictions.json',
         revisions: new Set(revisions),
@@ -337,7 +331,6 @@ describe('PathsModel', () => {
       {
         descendantStatuses: [2, 2, 2],
         hasChildren: true,
-        label: 'plots',
         parentPath: undefined,
         path: 'plots',
         revisions: new Set(revisions),
@@ -346,7 +339,6 @@ describe('PathsModel', () => {
       {
         descendantStatuses: [2, 2],
         hasChildren: true,
-        label: 'logs',
         parentPath: undefined,
         path: 'logs',
         revisions: new Set(revisions),
@@ -355,7 +347,6 @@ describe('PathsModel', () => {
       {
         descendantStatuses: [],
         hasChildren: false,
-        label: 'predictions.json',
         parentPath: undefined,
         path: 'predictions.json',
         revisions: new Set(revisions),
@@ -369,7 +360,6 @@ describe('PathsModel', () => {
       {
         descendantStatuses: [],
         hasChildren: false,
-        label: 'loss.tsv',
         parentPath: 'logs',
         path: logsLoss,
         revisions: new Set(revisions),
@@ -379,7 +369,6 @@ describe('PathsModel', () => {
       {
         descendantStatuses: [],
         hasChildren: false,
-        label: 'acc.tsv',
         parentPath: 'logs',
         path: logsAcc,
         revisions: new Set(revisions),
@@ -400,7 +389,6 @@ describe('PathsModel', () => {
       {
         descendantStatuses: [],
         hasChildren: true,
-        label: 'loss.tsv',
         parentPath: 'logs',
         path: logsLoss,
         revisions: new Set(revisions),
@@ -410,7 +398,6 @@ describe('PathsModel', () => {
       {
         descendantStatuses: [],
         hasChildren: true,
-        label: 'acc.tsv',
         parentPath: 'logs',
         path: logsAcc,
         revisions: new Set(revisions),

diff --git a/extension/src/plots/paths/model.ts b/extension/src/plots/paths/model.ts
@@ -77,7 +77,6 @@ export class PathsModel extends PathSelectionModel<PlotPath> {
       ...element,
       descendantStatuses: this.getTerminalNodeStatuses(element.path),
       hasChildren: this.getHasChildren(element, multiSourceEncoding),
-      label: element.label,
       status: this.status[element.path]
     }))
   }

diff --git a/extension/src/plots/paths/tree.test.ts b/extension/src/plots/paths/tree.test.ts
@@ -56,7 +56,11 @@ describe('PlotsPathsTree', () => {
       }
     ])
 
-    const children = plotsPathTree.getRepositoryChildren(__dirname, undefined)
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    const children = (plotsPathTree as any).getRepositoryChildren(
+      __dirname,
+      undefined
+    )
     expect(children).toStrictEqual([
       {
         collapsibleState: 0,

diff --git a/extension/src/plots/paths/tree.ts b/extension/src/plots/paths/tree.ts
@@ -1,4 +1,5 @@
-import { TreeItemCollapsibleState } from 'vscode'
+import { join } from 'path'
+import { TreeItem, TreeItemCollapsibleState } from 'vscode'
 import { EncodingType, isEncodingElement } from './collect'
 import {
   BasePathSelectionTree,
@@ -9,6 +10,7 @@ import { ResourceLocator } from '../../resourceLocator'
 import { RegisteredCommands } from '../../commands/external'
 import { EventName } from '../../telemetry/constants'
 import { InternalCommands } from '../../commands/internal'
+import { DecoratableTreeItemScheme, getDecoratableUri } from '../../tree'
 
 export class PlotsPathsTree extends BasePathSelectionTree<WorkspacePlots> {
   constructor(
@@ -32,7 +34,19 @@ export class PlotsPathsTree extends BasePathSelectionTree<WorkspacePlots> {
     )
   }
 
-  public getRepositoryChildren(dvcRoot: string, path: string | undefined) {
+  protected getBaseTreeItem({
+    dvcRoot,
+    path,
+    collapsibleState
+  }: PathSelectionItem) {
+    const resourceUri = getDecoratableUri(
+      join(dvcRoot, path),
+      DecoratableTreeItemScheme.PLOTS
+    )
+    return new TreeItem(resourceUri, collapsibleState)
+  }
+
+  protected getRepositoryChildren(dvcRoot: string, path: string | undefined) {
     return this.workspace
       .getRepository(dvcRoot)
       .getChildPaths(path)
@@ -56,7 +70,7 @@ export class PlotsPathsTree extends BasePathSelectionTree<WorkspacePlots> {
       })
   }
 
-  public getRepositoryStatuses(dvcRoot: string) {
+  protected getRepositoryStatuses(dvcRoot: string) {
     return this.workspace.getRepository(dvcRoot).getPathStatuses()
   }
 }