Add flexible plots legends to plots tree

extension/src/experiments/columns/model.ts

@@ -63,7 +63,22 @@ export class ColumnsModel extends PathSelectionModel<Column> {
     )
   }
 
-  public filterChildren(path?: string) {
+  public getChildren(path: string | undefined) {
+    return this.filterChildren(path).map(element => {
+      return {
+        ...element,
+        descendantStatuses: this.getTerminalNodeStatuses(element.path),
+        label: element.label,
+        status: this.status[element.path]
+      }
+    })
+  }
+
+  public hasNonDefaultColumns() {
+    return this.data.length > 1
+  }
+
+  private filterChildren(path?: string) {
     return this.data.filter(element =>
       path
         ? element.parentPath === path
@@ -73,10 +88,6 @@ export class ColumnsModel extends PathSelectionModel<Column> {
     )
   }
 
-  public hasNonDefaultColumns() {
-    return this.data.length > 1
-  }
-
   private async transformAndSetColumns(data: ExperimentsOutput) {
     const [columns, paramsFiles] = await Promise.all([
       collectColumns(data),

extension/src/experiments/columns/tree.ts

@@ -31,7 +31,10 @@ export class ExperimentsColumnsTree extends BasePathSelectionTree<WorkspaceExper
   }
 
   public getRepositoryChildren(dvcRoot: string, path: string) {
-    return this.workspace.getRepository(dvcRoot).getChildColumns(path)
+    return this.workspace
+      .getRepository(dvcRoot)
+      .getChildColumns(path)
+      .map(element => this.transformElement({ ...element, dvcRoot }))
   }
 
   public getRepositoryStatuses(dvcRoot: string) {

extension/src/path/selection/model.ts

@@ -44,17 +44,6 @@ export abstract class PathSelectionModel<
       .map(element => ({ ...element, selected: !!this.status[element.path] }))
   }
 
-  public getChildren(path: string | undefined) {
-    return this.filterChildren(path).map(element => {
-      return {
-        ...element,
-        descendantStatuses: this.getTerminalNodeStatuses(element.path),
-        label: element.label,
-        status: this.status[element.path]
-      }
-    })
-  }
-
   public toggleStatus(path: string) {
     const status = this.getNextStatus(path)
     this.status[path] = status
@@ -153,5 +142,7 @@ export abstract class PathSelectionModel<
     return this.persist(this.statusKey, this.status)
   }
 
-  abstract filterChildren(path?: string): T[]
+  abstract getChildren(
+    ...args: unknown[]
+  ): (T & { descendantStatuses: Status[]; label: string; status: Status })[]
 }

extension/src/path/selection/tree.ts

@@ -23,7 +23,7 @@ export type PathSelectionItem = {
   collapsibleState: TreeItemCollapsibleState
   label: string
   path: string
-  iconPath: Resource
+  iconPath: Resource | Uri
 }
 
 export abstract class BasePathSelectionTree<
@@ -35,8 +35,7 @@ export abstract class BasePathSelectionTree<
   public readonly onDidChangeTreeData: Event<PathSelectionItem | void>
 
   protected readonly workspace: T
-
-  private readonly resourceLocator: ResourceLocator
+  protected readonly resourceLocator: ResourceLocator
 
   private readonly view: TreeView<string | PathSelectionItem>
 
@@ -127,6 +126,33 @@ export abstract class BasePathSelectionTree<
     }${separator}${statuses.length}`
   }
 
+  protected transformElement(element: {
+    dvcRoot: string
+    descendantStatuses: Status[]
+    hasChildren: boolean
+    path: string
+    status: Status
+    label: string
+  }) {
+    const { dvcRoot, descendantStatuses, hasChildren, path, status, label } =
+      element
+
+    const description = this.getDescription(descendantStatuses, '/')
+    const iconPath = this.getIconPath(status)
+    const collapsibleState = hasChildren
+      ? TreeItemCollapsibleState.Collapsed
+      : TreeItemCollapsibleState.None
+
+    return {
+      collapsibleState,
+      description,
+      dvcRoot,
+      iconPath,
+      label,
+      path
+    } as PathSelectionItem
+  }
+
   private updateDescriptionOnChange() {
     this.dispose.track(
       this.onDidChangeTreeData(() => {
@@ -165,52 +191,22 @@ export abstract class BasePathSelectionTree<
     }
 
     if (this.isRoot(element)) {
-      return this.transformRepositoryChildren(element, undefined)
+      return this.getRepositoryChildren(element, undefined)
     }
 
     const { dvcRoot, path } = element
 
-    return this.transformRepositoryChildren(dvcRoot, path)
+    return this.getRepositoryChildren(dvcRoot, path)
   }
 
   private isRoot(element: string | PathSelectionItem): element is string {
     return typeof element === 'string'
   }
 
-  private transformRepositoryChildren(
-    dvcRoot: string,
-    path: string | undefined
-  ) {
-    return this.getRepositoryChildren(dvcRoot, path).map(element => {
-      const { descendantStatuses, hasChildren, path, status, label } = element
-
-      const description = this.getDescription(descendantStatuses, '/')
-      const iconPath = this.getIconPath(status)
-      const collapsibleState = hasChildren
-        ? TreeItemCollapsibleState.Collapsed
-        : TreeItemCollapsibleState.None
-
-      return {
-        collapsibleState,
-        description,
-        dvcRoot,
-        iconPath,
-        label,
-        path
-      } as PathSelectionItem
-    })
-  }
-
   abstract getRepositoryStatuses(dvcRoot: string): Status[]
 
   abstract getRepositoryChildren(
     dvcRoot: string,
     path: string | undefined
-  ): {
-    descendantStatuses: Status[]
-    hasChildren: boolean
-    label: string
-    path: string
-    status: Status
-  }[]
+  ): PathSelectionItem[]
 }

extension/src/plots/index.ts

@@ -4,7 +4,7 @@ import { PlotsData as TPlotsData } from './webview/contract'
 import { WebviewMessages } from './webview/messages'
 import { PlotsData } from './data'
 import { PlotsModel } from './model'
-import { collectScale } from './paths/collect'
+import { collectEncodingElements, collectScale } from './paths/collect'
 import { PathsModel } from './paths/model'
 import { BaseWebview } from '../webview'
 import { ViewKey } from '../webview/constants'
@@ -127,8 +127,14 @@ export class Plots extends BaseRepository<TPlotsData> {
     this.data.managedUpdate()
   }
 
-  public getChildPaths(path: string) {
-    return this.paths?.getChildren(path) || []
+  public getChildPaths(path: string | undefined) {
+    const multiSourceEncoding = this.plots?.getMultiSourceData() || {}
+
+    if (path && multiSourceEncoding[path]) {
+      return collectEncodingElements(path, multiSourceEncoding)
+    }
+
+    return this.paths?.getChildren(path, multiSourceEncoding) || []
   }
 
   public getPathStatuses() {

extension/src/plots/model/index.ts

@@ -306,6 +306,10 @@ export class PlotsModel extends ModelWithPersistence {
     return this.sectionCollapsed
   }
 
+  public getMultiSourceData() {
+    return this.multiSourceEncoding
+  }
+
   private removeStaleData() {
     return Promise.all([
       this.removeStaleBranches(),

extension/src/plots/multiSource/constants.ts

@@ -9,13 +9,7 @@ export const StrokeDash = [
 ] as const
 export type StrokeDashValue = typeof StrokeDash[number]
 
-export const Shape = [
-  'square',
-  'circle',
-  'triangle',
-  'diamond',
-  'cross'
-] as const
+export const Shape = ['square', 'circle', 'triangle', 'diamond'] as const
 export type ShapeValue = typeof Shape[number]
 
 export type Scale<T extends StrokeDashValue | ShapeValue> = {

extension/src/plots/paths/collect.test.ts

@@ -1,8 +1,14 @@
 import { join } from 'path'
 import { VisualizationSpec } from 'react-vega'
-import { collectPaths, collectTemplateOrder } from './collect'
+import {
+  collectEncodingElements,
+  collectPaths,
+  collectTemplateOrder,
+  EncodingType
+} from './collect'
 import { TemplatePlotGroup, PlotsType } from '../webview/contract'
 import plotsDiffFixture from '../../test/fixtures/plotsDiff/output'
+import { Shape, StrokeDash } from '../multiSource/constants'
 
 describe('collectPath', () => {
   it('should return the expected data from the test fixture', () => {
@@ -380,3 +386,55 @@ describe('collectTemplateOrder', () => {
     ])
   })
 })
+
+describe('collectEncodingElements', () => {
+  it('should return an empty array if there is no multi source encoding for a path', () => {
+    const elements = collectEncodingElements(__filename, {})
+    expect(elements).toStrictEqual([])
+  })
+
+  it('should collect encoding elements from multi source encoding', () => {
+    const elements = collectEncodingElements(__filename, {
+      [__filename]: {
+        shape: {
+          field: 'filename',
+          scale: { domain: ['X', 'Y'], range: [Shape[0], Shape[1]] }
+        },
+        strokeDash: {
+          field: 'field',
+          scale: {
+            domain: ['A', 'B', 'C'],
+            range: [StrokeDash[0], StrokeDash[1], StrokeDash[2]]
+          }
+        }
+      }
+    })
+    expect(elements).toStrictEqual([
+      {
+        label: 'A',
+        type: EncodingType.STROKE_DASH,
+        value: StrokeDash[0]
+      },
+      {
+        label: 'B',
+        type: EncodingType.STROKE_DASH,
+        value: StrokeDash[1]
+      },
+      {
+        label: 'C',
+        type: EncodingType.STROKE_DASH,
+        value: StrokeDash[2]
+      },
+      {
+        label: 'X',
+        type: EncodingType.SHAPE,
+        value: Shape[0]
+      },
+      {
+        label: 'Y',
+        type: EncodingType.SHAPE,
+        value: Shape[1]
+      }
+    ])
+  })
+})