TanStack cell selection

DataTable.tsx

"use client";

import {
  type ColumnDef,
  flexRender,
  getCoreRowModel,
  useReactTable,
  getPaginationRowModel,
  getSortedRowModel,
  getFilteredRowModel,
  getFacetedUniqueValues,
  type RowData,
} from "@tanstack/react-table";

import { cn } from "~/lib/utils";
import { DataTablePagination } from "./DataTablePagination";
import { DataTableFilter } from "./filters/DataTableFilter";
import { DebouncedInput } from "../inputs/DebouncedInput";
import { DataTableRowCount } from "./DataTableRowCount";
import { DataTableColumnHeader } from "./DataTableColumnHeader";
import { useQueryStringState } from "./hooks/useQueryStringState";
import { type ReactNode, useState } from "react";
import { Button } from "../ui/button";
import {
  DownloadIcon,
  FolderOpenIcon,
  MaximizeIcon,
  MinimizeIcon,
} from "lucide-react";
import { useCellSelection } from "./hooks/useCellSelection";
import { useCsvExport } from "./hooks/useCsvExport";
import { useRowVirtualizer } from "./hooks/useRowVirtualizer";
import { LoadingIcon } from "yet-another-react-lightbox";

declare module "@tanstack/react-table" {
  // eslint-disable-next-line @typescript-eslint/no-unused-vars
  interface ColumnMeta<TData extends RowData, TValue> {
    flex?: boolean;
  }
}

export type DataTableProps<TData, TValue> = {
  className?: string;
  columns: ColumnDef<TData, TValue>[];
  data: TData[];
  pagination: boolean;
  enableTopBar?: boolean;
  enableHeaderFilter?: boolean;
  isLoading?: boolean;
  maxCellLines?: number;
  children?: ReactNode;
};

export function DataTable<TData, TValue>({
  className,
  columns,
  data,
  children,
  pagination,
  enableTopBar = true,
  enableHeaderFilter = true,
  isLoading,
  maxCellLines = 3,
}: DataTableProps<TData, TValue>) {
  const [isFullscreen, setIsFullscreen] = useState(false);

  const {
    sorting,
    setSorting,
    columnFilters,
    setColumnFilters,
    globalFilter,
    setGlobalFilter,
  } = useQueryStringState();

  const table = useReactTable({
    data,
    columns,
    defaultColumn: {
      sortUndefined: "last",
      //size: "auto" as unknown as number,
      minSize: 10,
      maxSize: 500,
    },
    getCoreRowModel: getCoreRowModel(),
    ...(pagination && { getPaginationRowModel: getPaginationRowModel() }),
    onSortingChange: setSorting,
    getSortedRowModel: getSortedRowModel(),
    state: {
      sorting,
      columnFilters,
      globalFilter,
    },
    onColumnFiltersChange: setColumnFilters,
    onGlobalFilterChange: setGlobalFilter,
    globalFilterFn: "includesString",
    getColumnCanGlobalFilter: (column) => column.getCanFilter(),
    getFilteredRowModel: getFilteredRowModel(),
    getFacetedUniqueValues: getFacetedUniqueValues(),
  });

  const { rows } = table.getRowModel();
  const {
    tableContainerRef,
    tableRowRef,
    virtualRows,
    virtualHeight,
    scrollToRow,
  } = useRowVirtualizer({
    table,
  });

  const { isCellSelected, isRowSelected, isCellCopied, ...cellSelection } =
    useCellSelection({
      table,
      scrollToRow,
    });

  const { exportData } = useCsvExport({
    table,
  });

  const hasRows = virtualRows?.length > 0;

  return (
    <div
      className={cn(
        "flex min-h-48 flex-1 flex-col gap-2 overflow-hidden bg-background p-1",
        className,
        isFullscreen &&
          "height-[100dvh] width-[100dvw] fixed bottom-0 left-0 right-0 top-0 z-30 !m-0 max-h-[100dvh] p-2",
      )}
    >
      {enableTopBar && (
        <div className="flex items-center gap-2">
          {children}
          <DebouncedInput
            value={globalFilter ?? ""}
            onChange={(value) => setGlobalFilter(String(value))}
            className="font-lg border-blockmb-0 flex-1 border p-2 shadow"
            placeholder="Search all columns..."
          />
          <div className="ml-auto">
            <Button variant="ghost" size="icon" onClick={() => exportData()}>
              <DownloadIcon />
            </Button>
            <Button
              variant="ghost"
              size="icon"
              onClick={() => setIsFullscreen((fullscreen) => !fullscreen)}
            >
              {isFullscreen ? <MinimizeIcon /> : <MaximizeIcon />}
            </Button>
          </div>
        </div>
      )}
      <div className="flex flex-1 flex-col overflow-hidden rounded-md border">
        <div
          ref={tableContainerRef}
          className="flex flex-1 caption-bottom flex-col overflow-auto text-sm "
        >
          <div className="sticky top-0 z-10 flex">
            {table.getHeaderGroups().map((headerGroup) => (
              <div key={headerGroup.id} className="flex w-full">
                {headerGroup.headers.map((header) => {
                  return (
                    <div
                      key={header.id}
                      className="flex flex-col border-b bg-background p-1"
                      style={{
                        width: header.getSize(),
                        minWidth: header.getSize(),
                        flex: header.column.columnDef.meta?.flex
                          ? 1
                          : undefined,
                      }}
                    >
                      {header.isPlaceholder ? null : (
                        <>
                          <DataTableColumnHeader column={header.column}>
                            {flexRender(
                              header.column.columnDef.header,
                              header.getContext(),
                            )}
                          </DataTableColumnHeader>
                          {(enableHeaderFilter || isFullscreen) &&
                          header.column.getCanFilter() ? (
                            <DataTableFilter
                              column={header.column}
                              table={table}
                            />
                          ) : null}
                        </>
                      )}
                    </div>
                  );
                })}
              </div>
            ))}
          </div>
          <div className="relative flex flex-1 overflow-visible">
            <div
              className="relative flex flex-1"
              style={{
                height: virtualHeight,
              }}
              tabIndex={-1}
              onKeyDown={cellSelection.handleCellsKeyDown}
            >
              {hasRows &&
                virtualRows.map((virtualRow) => {
                  const row = rows[virtualRow.index]!;

                  return (
                    <div
                      key={row.id}
                      ref={tableRowRef}
                      data-state={row.getIsSelected() && "selected"}
                      data-index={virtualRow.index}
                      className="w-100 absolute left-0 right-0 flex border-b border-b-background transition-colors hover:bg-muted data-[state=selected]:bg-muted"
                      style={{
                        transform: `translateY(${virtualRow.start}px)`,
                      }}
                    >
                      {row.getVisibleCells().map((cell) => (
                        <div
                          onMouseDown={(e) =>
                            cellSelection.handleCellMouseDown(e, cell)
                          }
                          onMouseUp={(e) =>
                            cellSelection.handleCellMouseUp(e, cell)
                          }
                          onMouseOver={(e) =>
                            cellSelection.handleCellMouseOver(e, cell)
                          }
                          key={cell.id}
                          className={cn(
                            "relative flex select-none overflow-hidden border-x border-b border-x-transparent bg-background p-2 align-middle transition-colors before:transition-colors focus:outline-none [&:has([role=checkbox])]:pr-0 [&>[role=checkbox]]:translate-y-[2px]",
                            isCellSelected(cell) &&
                              "bg-primary/10 before:pointer-events-none before:absolute before:bottom-0 before:left-0 before:right-0 before:top-0 before:border before:border-primary hover:before:bg-primary/10",
                            isCellCopied(cell) &&
                              "bg-copy duration-300 before:border-copy-border before:duration-300",
                          )}
                          style={{
                            width: cell.column.getSize(),
                            minWidth: cell.column.getSize(),
                            flex: cell.column.columnDef.meta?.flex
                              ? 1
                              : undefined,
                          }}
                        >
                          <div
                            className="overflow-hidden"
                            style={{
                              display: "-webkit-box",
                              WebkitLineClamp: !isRowSelected(cell.row.id)
                                ? maxCellLines
                                : undefined,
                              WebkitBoxOrient: "vertical",
                              textOverflow: "clip",
                              lineHeight: "1.4rem",
                            }}
                          >
                            {flexRender(
                              cell.column.columnDef.cell,
                              cell.getContext(),
                            )}
                          </div>
                        </div>
                      ))}
                    </div>
                  );
                })}
              {!hasRows && !isLoading && (
                <div className="flex h-24 flex-1 items-center justify-center gap-2">
                  <FolderOpenIcon className="text-primary/50" />
                  No results.
                </div>
              )}
              {!hasRows && isLoading && (
                <div className="flex h-24 flex-1 items-center justify-center gap-2">
                  <LoadingIcon className="animate-spin text-primary/50" />
                  Loading...
                </div>
              )}
            </div>
          </div>
        </div>
        <div className="flex justify-between border-t p-2 pl-4">
          <DataTableRowCount table={table} />
          {pagination && <DataTablePagination table={table} />}
          <Button
            variant="ghost"
            size={null}
            onClick={() => setIsFullscreen((fullscreen) => !fullscreen)}
          >
            {isFullscreen ? (
              <MinimizeIcon size={15} />
            ) : (
              <MaximizeIcon size={15} />
            )}
          </Button>
        </div>
      </div>
    </div>
  );
}

useCellSelection.ts

import type { Cell, Table } from "@tanstack/react-table";
import { useState } from "react";
import { useCopyToClipboard } from "~/hooks/useCopyToClipboard";
import { useMouseOut } from "~/hooks/useMouseOut";

export type UseCellSelectionProps = {
  table: Table<any>;
  scrollToRow?: (index: number) => void;
};

export type SelectedCell = {
  rowId: string;
  columnId: string;
  cellId: string;
};

export const useCellSelection = ({
  table,
  scrollToRow,
}: UseCellSelectionProps) => {
  const [selectedCells, setSelectedCells] = useState<SelectedCell[]>([]);
  const [copiedCells, setCopiedCells] = useState<SelectedCell[]>([]);

  const [selectedStartCell, setSelectedStartCell] =
    useState<SelectedCell | null>(null);
  const [isMouseDown, setIsMouseDown] = useState(false);

  const [_copiedText, copy] = useCopyToClipboard();
  const handleCopy = () => {
    // eslint-disable-next-line @typescript-eslint/no-floating-promises
    copy(getCellValues(table, selectedCells));

    setCopiedCells(selectedCells);
    setTimeout(() => {
      setCopiedCells([]);
    }, 500);
  };

  const handleCellsKeyDown = (e: React.KeyboardEvent<HTMLElement>) => {
    switch (e.key) {
      case "c": {
        if (e.metaKey || e.ctrlKey) {
          handleCopy();
        }
        break;
      }
      case "ArrowDown": {
        e.preventDefault();
        navigateDown();
        break;
      }
      case "ArrowUp": {
        e.preventDefault();
        navigateUp();
        break;
      }
      case "ArrowLeft": {
        e.preventDefault();
        navigateLeft();
        break;
      }
      case "ArrowRight": {
        e.preventDefault();
        navigateRight();
        break;
      }
      case "Home": {
        e.preventDefault();
        navigateHome();
        break;
      }
      case "End": {
        e.preventDefault();
        navigateEnd();
        break;
      }
    }
  };

  useMouseOut(() => {
    setIsMouseDown(false);
  });

  const navigateHome = () => {
    const firstCell = table.getRowModel().rows[0]?.getAllCells()[0];
    if (!firstCell) {
      return;
    }

    setSelectedCells([getCellSelectionData(firstCell)]);
    scrollToRow?.(0);
  };

  const navigateEnd = () => {
    const lastRow =
      table.getRowModel().rows[table.getRowModel().rows.length - 1];
    const lastCell = lastRow?.getAllCells()[lastRow.getAllCells().length - 1];
    if (!lastCell) {
      return;
    }

    setSelectedCells([getCellSelectionData(lastCell)]);
    scrollToRow?.(table.getRowModel().rows.length);
  };

  const navigateUp = () => {
    const selectedCell = selectedCells[selectedCells.length - 1];
    if (!selectedCell) {
      return;
    }

    const selectedRowIndex = table
      .getRowModel()
      .rows.findIndex((row) => row.id === selectedCell.rowId);
    const nextRowIndex = selectedRowIndex - 1;
    const previousRow = table.getRowModel().rows[nextRowIndex];
    if (previousRow) {
      setSelectedCells([
        getCellSelectionData(
          previousRow
            .getAllCells()
            .find((c) => c.column.id === selectedCell.columnId)!,
        ),
      ]);
      scrollToRow?.(nextRowIndex);
    }
  };

  const navigateDown = () => {
    const selectedCell = selectedCells[selectedCells.length - 1];
    if (!selectedCell) {
      return;
    }

    const selectedRowIndex = table
      .getRowModel()
      .rows.findIndex((row) => row.id === selectedCell.rowId);
    const nextRowIndex = selectedRowIndex + 1;
    const nextRow = table.getRowModel().rows[nextRowIndex];
    if (nextRow) {
      setSelectedCells([
        getCellSelectionData(
          nextRow
            .getAllCells()
            .find((c) => c.column.id === selectedCell.columnId)!,
        ),
      ]);
      scrollToRow?.(nextRowIndex);
    }
  };

  const navigateLeft = () => {
    const selectedCell = selectedCells[selectedCells.length - 1];
    if (!selectedCell) {
      return;
    }

    const selectedRow = table.getRow(selectedCell.rowId);
    const selectedColumnIndex = selectedRow
      .getAllCells()
      .findIndex((c) => c.id === selectedCell.cellId);
    const previousCell = selectedRow.getAllCells()[selectedColumnIndex - 1];
    if (previousCell) {
      setSelectedCells([getCellSelectionData(previousCell)]);
    }
  };

  const navigateRight = () => {
    const selectedCell = selectedCells[selectedCells.length - 1];
    if (!selectedCell) {
      return;
    }

    const selectedRow = table.getRow(selectedCell.rowId);
    const selectedColumnIndex = selectedRow
      .getAllCells()
      .findIndex((c) => c.id === selectedCell.cellId);
    const nextCell = selectedRow.getAllCells()[selectedColumnIndex + 1];
    if (nextCell) {
      setSelectedCells([getCellSelectionData(nextCell)]);
    }
  };

  const isRowSelected = (rowId: string) =>
    selectedCells.find((c) => c.rowId === rowId) !== undefined;

  const isCellSelected = (cell: Cell<any, any>) =>
    selectedCells.find((c) => c.cellId === cell.id) !== undefined;

  const isCellCopied = (cell: Cell<any, any>) =>
    copiedCells.find((c) => c.cellId === cell.id) !== undefined;

  const updateRangeSelection = (cell: Cell<any, any>) => {
    if (!selectedStartCell) {
      return;
    }

    const selectedCellsInRange = getCellsBetween(
      table,
      selectedStartCell,
      getCellSelectionData(cell),
    ) as SelectedCell[];

    setSelectedCells((prev) => {
      const startIndex = prev.findIndex(
        (c) => c.cellId === selectedStartCell.cellId,
      );
      const prevSelectedCells = prev.slice(0, startIndex);
      const newCellSelection = selectedCellsInRange.filter(
        (c) => c.cellId !== selectedStartCell.cellId,
      );

      return [...prevSelectedCells, selectedStartCell, ...newCellSelection];
    });
  };

  const handleCellMouseDown = (
    e: React.MouseEvent<HTMLElement>,
    cell: Cell<any, any>,
  ) => {
    if (!e.ctrlKey && !e.shiftKey) {
      setSelectedCells([getCellSelectionData(cell)]);
      if (!isMouseDown) {
        setSelectedStartCell(getCellSelectionData(cell));
      }
    }

    if (e.ctrlKey) {
      setSelectedCells((prev) =>
        prev.find((c) => c.cellId === cell.id) !== undefined
          ? prev.filter(({ cellId }) => cellId !== cell.id)
          : [...prev, getCellSelectionData(cell)],
      );
      if (!isMouseDown) {
        setSelectedStartCell(getCellSelectionData(cell));
      }
    }

    if (e.shiftKey) {
      updateRangeSelection(cell);
    }

    setIsMouseDown(true);
  };

  const handleCellMouseUp = (
    e: React.MouseEvent<HTMLElement>,
    _cell: Cell<any, any>,
  ) => {
    if (!e.shiftKey) {
    }

    setIsMouseDown(false);
  };

  const handleCellMouseOver = (
    e: React.MouseEvent<HTMLElement>,
    cell: Cell<any, any>,
  ) => {
    if (e.buttons !== 1) return;

    if (isMouseDown) {
      updateRangeSelection(cell);
    }
  };

  return {
    handleCellMouseDown,
    handleCellMouseUp,
    handleCellMouseOver,
    handleCellsKeyDown,
    isCellSelected,
    isRowSelected,
    isCellCopied,
  };
};

type SelectedCellRowMap = Record<string, SelectedCell[]>;
const getCellValues = (table: Table<any>, cells: SelectedCell[]) => {
  // reduce cells into arrays of rows
  const rows = cells.reduce(
    (acc: SelectedCellRowMap, cellIds: SelectedCell) => {
      const cellsForRow = acc[cellIds.rowId] ?? [];
      return {
        ...acc,
        [cellIds.rowId]: [...cellsForRow, cellIds],
      };
    },
    {} as SelectedCellRowMap,
  );

  return Object.keys(rows)
    .map((rowId) => {
      const selectedCells = rows[rowId]!;
      const row = table.getRow(rowId);

      const cellValues = [];
      for (const cell of row.getAllCells()) {
        if (selectedCells.find((c) => c.cellId === cell.id)) {
          cellValues.push(cell?.getValue());
        }
      }

      return cellValues.join("\t");
    })
    .join("\n");
};

const getCellSelectionData = (cell: Cell<any, any>) => ({
  rowId: cell.row.id,
  columnId: cell.column.id,
  cellId: cell.id,
});

const getSelectedCellTableData = (table: Table<any>, cell: SelectedCell) => {
  const row = table.getRow(cell.rowId);
  return row.getAllCells().find((c) => c.id === cell.cellId);
};

const getCellsBetween = (
  table: Table<any>,
  cell1: SelectedCell,
  cell2: SelectedCell,
) => {
  const cell1Data = getSelectedCellTableData(table, cell1);
  const cell2Data = getSelectedCellTableData(table, cell2);
  if (!cell1Data || !cell2Data) return [];

  const rows = table.getRowModel().rows;

  const cell1RowIndex = rows.findIndex(({ id }) => id === cell1Data.row.id);
  const cell2RowIndex = rows.findIndex(({ id }) => id === cell2Data.row.id);

  const cell1ColumnIndex = cell1Data.column.getIndex();
  const cell2ColumnIndex = cell2Data.column.getIndex();

  const selectedRows = rows.slice(
    Math.min(cell1RowIndex, cell2RowIndex),
    Math.max(cell1RowIndex, cell2RowIndex) + 1,
  );

  const columns = table
    .getAllColumns()
    .slice(
      Math.min(cell1ColumnIndex, cell2ColumnIndex),
      Math.max(cell1ColumnIndex, cell2ColumnIndex) + 1,
    );

  return selectedRows.flatMap((row) =>
    columns.map((column) => {
      const tableCell = row
        .getAllCells()
        .find((cell) => cell.column.id === column.id);
      if (!tableCell) return null;
      return getCellSelectionData(tableCell);
    }),
  );
};

Great find! I only changed getAllColumns to getAllLeafColumns, so it continues to work with nested neaders https://gist.github.com/joshkay/fc8bab0561583dd48fecce93022fc7a2#file-usecellselection-ts-L343