frontend/lib/src/components/widgets/DataFrame/hooks/useSelectionHandler.ts - develop - vshmidt/streamlit

frontend/lib/src/components/widgets/DataFrame/hooks/useSelectionHandler.ts
235 строк8 KB
Lukas Masuch
Split `Arrow.proto` into `Table.proto` and `Dataframe.proto` (#13768)
04 фев 2026, 22:22
04 фев 2026, 22:22dc9d8fe
/**
 * Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2026)
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import { useCallback, useState } from "react"

import { CompactSelection, GridSelection } from "@glideapps/glide-data-grid"
import { isEqual } from "lodash-es"

import { Dataframe as DataframeProto } from "@streamlit/protobuf"

import { BaseColumn } from "~lib/components/widgets/DataFrame/columns"

export type SelectionHandlerReturn = {
  // The current selection state
  gridSelection: GridSelection
  // True, if row selection is activated
  isRowSelectionActivated: boolean
  // True, if multi row selection is activated
  isMultiRowSelectionActivated: boolean
  // True, if column selection is activated
  isColumnSelectionActivated: boolean
  // True, if multi column selections is activated
  isMultiColumnSelectionActivated: boolean
  // True, if cell selection is activated
  isCellSelectionActivated: boolean
  // True, if multi cell selection is activated
  isMultiCellSelectionActivated: boolean
  // True, if at least one row is selected
  isRowSelected: boolean
  // True, if at least one column is selected
  isColumnSelected: boolean
  // True, if at least one cell is selected
  isCellSelected: boolean
  // Callback to clear selections
  clearSelection: (keepRows?: boolean, keepColumns?: boolean) => void
  // Callback to process selection changes from the grid
  processSelectionChange: (newSelection: GridSelection) => void
}

/**
 * Custom hook that handles all selection capabilities for the interactive data table.
 *
 * @param element - The Arrow proto message
 * @param isEmptyTable - Whether the table is empty
 * @param isDisabled - Whether the table is disabled
 * @param columns - The columns of the table.
 * @param syncSelectionState - The callback to sync the selection state
 *
 * @returns the selection handler return object
 */
function useSelectionHandler(
  element: DataframeProto,
  isEmptyTable: boolean,
  isDisabled: boolean,
  columns: BaseColumn[],
  syncSelectionState: (
    newSelection: GridSelection,
    syncCellSelections: boolean
  ) => void
): SelectionHandlerReturn {
  const [gridSelection, setGridSelection] = useState<GridSelection>({
    columns: CompactSelection.empty(),
    rows: CompactSelection.empty(),
    current: undefined,
  })

  const isRowSelectionActivated =
    !isEmptyTable &&
    !isDisabled &&
    (element.selectionMode.includes(DataframeProto.SelectionMode.MULTI_ROW) ||
      element.selectionMode.includes(DataframeProto.SelectionMode.SINGLE_ROW))
  const isMultiRowSelectionActivated =
    isRowSelectionActivated &&
    element.selectionMode.includes(DataframeProto.SelectionMode.MULTI_ROW)

  const isColumnSelectionActivated =
    !isEmptyTable &&
    !isDisabled &&
    (element.selectionMode.includes(
      DataframeProto.SelectionMode.SINGLE_COLUMN
    ) ||
      element.selectionMode.includes(
        DataframeProto.SelectionMode.MULTI_COLUMN
      ))
  const isMultiColumnSelectionActivated =
    isColumnSelectionActivated &&
    element.selectionMode.includes(DataframeProto.SelectionMode.MULTI_COLUMN)

  const isCellSelectionActivated =
    !isEmptyTable &&
    !isDisabled &&
    (element.selectionMode.includes(
      DataframeProto.SelectionMode.SINGLE_CELL
    ) ||
      element.selectionMode.includes(DataframeProto.SelectionMode.MULTI_CELL))

  const isMultiCellSelectionActivated =
    isCellSelectionActivated &&
    element.selectionMode.includes(DataframeProto.SelectionMode.MULTI_CELL)

  const isRowSelected = gridSelection.rows.length > 0
  const isColumnSelected = gridSelection.columns.length > 0
  const isCellSelected = gridSelection.current !== undefined

  /**
   * This callback is used to process selection changes and - if activated -
   * trigger a sync of the state with the widget state
   */
  const processSelectionChange = useCallback(
    // eslint-disable-next-line react-hooks/preserve-manual-memoization -- TODO: Update to match React best practices
    (newSelection: GridSelection) => {
      const rowSelectionChanged = !isEqual(
        newSelection.rows.toArray(),
        gridSelection.rows.toArray()
      )

      const columnSelectionChanged = !isEqual(
        newSelection.columns.toArray(),
        gridSelection.columns.toArray()
      )

      const cellSelectionChanged = !isEqual(
        newSelection.current,
        gridSelection.current
      )

      // A flag to determine if the selection should be synced with the widget state
      const syncSelection =
        (isRowSelectionActivated && rowSelectionChanged) ||
        (isColumnSelectionActivated && columnSelectionChanged) ||
        (isCellSelectionActivated && cellSelectionChanged)

      let updatedSelection = newSelection

      if (columnSelectionChanged && updatedSelection.columns.length >= 0) {
        // Remove all index columns from the column selection
        // We don't want to allow selection of index columns.
        let cleanedColumns = updatedSelection.columns
        columns.forEach((column, idx) => {
          if (column.isIndex) {
            cleanedColumns = cleanedColumns.remove(idx)
          }
        })
        if (cleanedColumns.length < updatedSelection.columns.length) {
          updatedSelection = {
            ...updatedSelection,
            columns: cleanedColumns,
          }
        }
      }

      // Update the UI with the final selection state
      setGridSelection(updatedSelection)

      if (syncSelection) {
        // Sync this selection with the widget state / backend
        syncSelectionState(updatedSelection, isCellSelectionActivated)
      }
    },
    [
      gridSelection,
      isRowSelectionActivated,
      isColumnSelectionActivated,
      isCellSelectionActivated,
      syncSelectionState,
      columns,
    ]
  )

  /**
   * This callback is used to selections (row/column/cell)
   * and sync the state with the widget state if column or row selections
   * are activated and the selection has changed.
   *
   * @param keepRows - Whether to keep the row selection (default: false)
   * @param keepColumns - Whether to keep the column selection (default: false)
   */
  const clearSelection = useCallback(
    (keepRows = false, keepColumns = false) => {
      const emptySelection: GridSelection = {
        columns: keepColumns
          ? gridSelection.columns
          : CompactSelection.empty(),
        rows: keepRows ? gridSelection.rows : CompactSelection.empty(),
        current: undefined,
      }
      setGridSelection(emptySelection)

      if (
        (!keepRows && isRowSelectionActivated) ||
        (!keepColumns && isColumnSelectionActivated) ||
        isCellSelectionActivated
      ) {
        syncSelectionState(emptySelection, isCellSelectionActivated)
      }
    },
    [
      gridSelection,
      isRowSelectionActivated,
      isColumnSelectionActivated,
      isCellSelectionActivated,
      syncSelectionState,
    ]
  )

  return {
    gridSelection,
    isRowSelectionActivated,
    isMultiRowSelectionActivated,
    isColumnSelectionActivated,
    isMultiColumnSelectionActivated,
    isCellSelectionActivated,
    isMultiCellSelectionActivated,
    isRowSelected,
    isColumnSelected,
    isCellSelected,
    clearSelection,
    processSelectionChange,
  }
}

export default useSelectionHandler