Refactor experiment table header context menu (#4407)

webview/src/experiments/components/table/Table.test.tsx

@@ -15,7 +15,7 @@ import React from 'react'
 import tableDataFixture from 'dvc/src/test/fixtures/expShow/base/tableData'
 import { EXPERIMENT_WORKSPACE_ID } from 'dvc/src/cli/dvc/contract'
 import styles from './styles.module.scss'
-import { SortOrder } from './header/ContextMenuContent'
+import { SortOrder } from './header/util'
 import { ExperimentsTable } from '../Experiments'
 import { vsCodeApi } from '../../../shared/api'
 import {

webview/src/experiments/components/table/header/ContextMenuContent.tsx

@@ -1,35 +1,21 @@
 import { ColumnType, Experiment } from 'dvc/src/experiments/webview/contract'
 import { MessageFromWebviewType } from 'dvc/src/webview/contract'
-import { VSCodeDivider } from '@vscode/webview-ui-toolkit/react'
 import React, { useMemo } from 'react'
 import { Header } from '@tanstack/react-table'
 import { useSelector } from 'react-redux'
 import { SortDefinition } from 'dvc/src/experiments/model/sortBy'
+import { SortOrder, getSortDetails, isFromExperimentColumn } from './util'
 import { MessagesMenu } from '../../../../shared/components/messagesMenu/MessagesMenu'
 import { MessagesMenuOptionProps } from '../../../../shared/components/messagesMenu/MessagesMenuOption'
 import { ExperimentsState } from '../../../store'
 import { ColumnWithGroup } from '../../../util/buildColumns'
 
-export enum SortOrder {
-  ASCENDING = 'Sort Ascending',
-  DESCENDING = 'Sort Descending',
-  NONE = 'Remove Sort'
-}
-
-const possibleOrders = {
-  false: SortOrder.ASCENDING,
-  true: SortOrder.DESCENDING,
-  undefined: SortOrder.NONE
-} as const
-
-const isFromExperimentColumn = (header: Header<Experiment, unknown>) =>
-  header.column.id === 'id' || header.column.id.startsWith('id_placeholder')
-
 const sortOption = (
   label: SortOrder,
   currentSort: SortOrder,
   columnId: string,
-  isSortable: boolean
+  isSortable: boolean,
+  divider?: boolean
 ) => {
   const sortOrder = currentSort
   const disabled = !isSortable || sortOrder === label
@@ -53,6 +39,7 @@ const sortOption = (
 
   return {
     disabled,
+    divider,
     id: label,
     label,
     message
@@ -63,36 +50,14 @@ interface HeaderMenuProps {
   header: Header<Experiment, unknown>
 }
 
-const getSortOptions = (
-  header: Header<Experiment, unknown>,
-  sorts: SortDefinition[]
-) => {
-  const isNotExperiments = !isFromExperimentColumn(header)
-  const isSortable = isNotExperiments && header.column.columns.length <= 1
-  const baseColumn =
-    header.headerGroup.headers.find(
-      h => h.column.id === header.placeholderId
-    ) || header.column
-  const sort = sorts.find(sort => sort.path === baseColumn.id)
-
-  // eslint-disable-next-line @typescript-eslint/restrict-template-expressions
-  const sortOrder: SortOrder = possibleOrders[`${sort?.descending}`]
-
-  const sortOptions = [
-    sortOption(SortOrder.ASCENDING, sortOrder, baseColumn.id, isSortable),
-    sortOption(SortOrder.DESCENDING, sortOrder, baseColumn.id, isSortable),
-    sortOption(SortOrder.NONE, sortOrder, baseColumn.id, isSortable)
-  ]
-
-  return { isSortable, sortOptions, sortOrder }
-}
-
-export const getMenuOptions = (
+const getMenuOptions = (
   header: Header<Experiment, unknown>,
   sorts: SortDefinition[]
-) => {
+): MessagesMenuOptionProps[] => {
   const leafColumn = header.column
-  const menuOptions: MessagesMenuOptionProps[] = [
+  const { id, isSortable, sortOrder } = getSortDetails(header, sorts)
+
+  return [
     {
       disabled: isFromExperimentColumn(header),
       id: 'hide',
@@ -143,30 +108,19 @@ export const getMenuOptions = (
       message: {
         type: MessageFromWebviewType.SELECT_FIRST_COLUMNS
       }
-    }
+    },
+    sortOption(SortOrder.ASCENDING, sortOrder, id, isSortable, true),
+    sortOption(SortOrder.DESCENDING, sortOrder, id, isSortable),
+    sortOption(SortOrder.NONE, sortOrder, id, isSortable)
   ]
-
-  const { isSortable, sortOptions, sortOrder } = getSortOptions(header, sorts)
-
-  return { isSortable, menuOptions, sortOptions, sortOrder }
 }
 
 export const ContextMenuContent: React.FC<HeaderMenuProps> = ({ header }) => {
   const { sorts } = useSelector((state: ExperimentsState) => state.tableData)
 
-  const { menuOptions, sortOptions } = useMemo(() => {
+  const menuOptions = useMemo(() => {
     return getMenuOptions(header, sorts)
   }, [header, sorts])
 
-  return (
-    <div>
-      <MessagesMenu options={menuOptions} />
-      {sortOptions.length > 0 && (
-        <>
-          <VSCodeDivider />
-          <MessagesMenu options={sortOptions} />
-        </>
-      )}
-    </div>
-  )
+  return <MessagesMenu options={menuOptions} />
 }

webview/src/experiments/components/table/header/MergeHeaderGroups.tsx

@@ -2,7 +2,7 @@ import React from 'react'
 import cx from 'classnames'
 import { Experiment } from 'dvc/src/experiments/webview/contract'
 import { HeaderGroup, Header } from '@tanstack/react-table'
-import { TableHeader } from './TableHeader'
+import { TableHeaderCell } from './TableHeaderCell'
 import styles from '../styles.module.scss'
 import { DragFunction } from '../../../../shared/components/dragDrop/Draggable'
 
@@ -30,7 +30,7 @@ export const MergedHeaderGroups: React.FC<{
   return (
     <tr className={cx(styles.experimentsTr, styles.headRow)}>
       {headerGroup.headers.map((header: Header<Experiment, unknown>) => (
-        <TableHeader
+        <TableHeaderCell
           setExpColumnNeedsShadow={setExpColumnNeedsShadow}
           key={header.id}
           header={header}

webview/src/experiments/components/table/header/TableHeaderCell.tsx

@@ -5,7 +5,8 @@ import { Header } from '@tanstack/react-table'
 import cx from 'classnames'
 import { useInView } from 'react-intersection-observer'
 import { TableHeaderCellContents } from './TableHeaderCellContents'
-import { ContextMenuContent, getMenuOptions } from './ContextMenuContent'
+import { ContextMenuContent } from './ContextMenuContent'
+import { getSortDetails } from './util'
 import styles from '../styles.module.scss'
 import { isExperimentColumn, isFirstLevelHeader } from '../../../util/columns'
 import { ExperimentsState } from '../../../store'
@@ -80,7 +81,6 @@ const WithExpColumnNeedsShadowUpdates: React.FC<{
 
 export const TableHeaderCell: React.FC<{
   header: Header<Experiment, unknown>
-  hasFilter: boolean
   onDragEnter: DragFunction
   onDragEnd: DragFunction
   onDragStart: DragFunction
@@ -91,7 +91,6 @@ export const TableHeaderCell: React.FC<{
   onlyOneLine?: boolean
 }> = ({
   header,
-  hasFilter,
   onDragEnter,
   onDragEnd,
   onDragStart,
@@ -111,13 +110,17 @@ export const TableHeaderCell: React.FC<{
   const headerDropTargetId = useSelector(
     (state: ExperimentsState) => state.headerDropTarget
   )
-  const { sorts } = useSelector((state: ExperimentsState) => state.tableData)
+  const { filters, sorts } = useSelector(
+    (state: ExperimentsState) => state.tableData
+  )
 
   const { isSortable, sortOrder } = useMemo(() => {
-    return getMenuOptions(header, sorts)
+    return getSortDetails(header, sorts)
   }, [header, sorts])
   const isDraggable = !isPlaceholder && !isExperimentColumn(id)
 
+  const hasFilter = !!(header.id && filters.includes(header.id))
+
   const canResize = getCanResize() && !isPlaceholder
   const resizerHeight = calcResizerHeight(header)
 

webview/src/experiments/components/table/header/TableHeaderCellContents.tsx

@@ -2,8 +2,8 @@ import React, { useEffect } from 'react'
 import cx from 'classnames'
 import { Experiment } from 'dvc/src/experiments/webview/contract'
 import { flexRender, Header } from '@tanstack/react-table'
-import { SortOrder } from './ContextMenuContent'
 import { ColumnResizer, ResizerHeight } from './ColumnResizer'
+import { SortOrder } from './util'
 import styles from '../styles.module.scss'
 import {
   Draggable,

webview/src/experiments/components/table/header/util.ts

@@ -0,0 +1,36 @@
+import { Header } from '@tanstack/react-table'
+import { SortDefinition } from 'dvc/src/experiments/model/sortBy'
+import { Experiment } from 'dvc/src/experiments/webview/contract'
+
+export enum SortOrder {
+  ASCENDING = 'Sort Ascending',
+  DESCENDING = 'Sort Descending',
+  NONE = 'Remove Sort'
+}
+
+const possibleOrders = {
+  false: SortOrder.ASCENDING,
+  true: SortOrder.DESCENDING,
+  undefined: SortOrder.NONE
+} as const
+
+export const isFromExperimentColumn = (header: Header<Experiment, unknown>) =>
+  header.column.id === 'id' || header.column.id.startsWith('id_placeholder')
+
+export const getSortDetails = (
+  header: Header<Experiment, unknown>,
+  sorts: SortDefinition[]
+): { id: string; isSortable: boolean; sortOrder: SortOrder } => {
+  const isNotExperiments = !isFromExperimentColumn(header)
+  const isSortable = isNotExperiments && header.column.columns.length <= 1
+  const baseColumn =
+    header.headerGroup.headers.find(
+      h => h.column.id === header.placeholderId
+    ) || header.column
+  const sort = sorts.find(sort => sort.path === baseColumn.id)
+
+  // eslint-disable-next-line @typescript-eslint/restrict-template-expressions
+  const sortOrder: SortOrder = possibleOrders[`${sort?.descending}`]
+
+  return { id: baseColumn.id, isSortable, sortOrder }
+}