feat: heatmap columns and full table in mo.ui.table

frontend/src/components/data-table/TableActions.tsx

@@ -2,10 +2,16 @@
 import React from "react";
 import { Tooltip } from "../ui/tooltip";
 import { Button } from "../ui/button";
-import { SearchIcon } from "lucide-react";
+import { PaletteIcon, SearchIcon, Settings } from "lucide-react";
 import { DataTablePagination } from "./pagination";
 import { DownloadAs, type DownloadActionProps } from "./download-actions";
 import type { Table, RowSelectionState } from "@tanstack/react-table";
+import {
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuItem,
+  DropdownMenuTrigger,
+} from "../ui/dropdown-menu";
 
 interface TableActionsProps<TData> {
   enableSearch: boolean;
@@ -44,7 +50,7 @@ export const TableActions = <TData,>({
           </Button>
         </Tooltip>
       )}
-      {pagination ? (
+      {pagination && (
         <DataTablePagination
           selection={selection}
           onSelectAllRowsChange={
@@ -64,10 +70,25 @@ export const TableActions = <TData,>({
           }
           table={table}
         />
-      ) : (
-        <div />
       )}
-      {downloadAs && <DownloadAs downloadAs={downloadAs} />}
+      <div className="flex items-center">
+        {downloadAs && <DownloadAs downloadAs={downloadAs} />}
+        {table.toggleGlobalHeatmap && (
+          <DropdownMenu>
+            <DropdownMenuTrigger asChild={true}>
+              <Button variant="text" size="xs" className="mb-0">
+                <Settings className="w-4 h-4 text-muted-foreground" />
+              </Button>
+            </DropdownMenuTrigger>
+            <DropdownMenuContent align="end">
+              <DropdownMenuItem onClick={() => table.toggleGlobalHeatmap?.()}>
+                <PaletteIcon className="mr-2 h-3.5 w-3.5 text-muted-foreground/70" />
+                {table.getGlobalHeatmap?.() ? "Disable" : "Enable"} heatmap
+              </DropdownMenuItem>
+            </DropdownMenuContent>
+          </DropdownMenu>
+        )}
+      </div>
     </div>
   );
 };

frontend/src/components/data-table/__test__/cell-heatmap.test.ts

@@ -0,0 +1,99 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { describe, it, expect, vi } from "vitest";
+import { CellHeatmapFeature } from "../cell-heatmap/feature";
+import {
+  type Column,
+  createTable,
+  getCoreRowModel,
+} from "@tanstack/react-table";
+
+describe("CellHeatmapFeature", () => {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  let state: any = {
+    columnHeatmap: {
+      value: true,
+    },
+    cachedMaxValue: null,
+    cachedMinValue: null,
+  };
+
+  const mockTable = createTable({
+    _features: [CellHeatmapFeature],
+    data: [
+      { id: 1, value: 10 },
+      { id: 2, value: 20 },
+      { id: 3, value: 30 },
+    ],
+    state: state,
+    onStateChange: (updater) => {
+      state = typeof updater === "function" ? updater(state) : updater;
+    },
+    columns: [
+      { id: "id", accessorKey: "id" },
+      { id: "value", accessorKey: "value" },
+    ],
+    getCoreRowModel: getCoreRowModel(),
+    renderFallbackValue: null,
+  });
+
+  it("should initialize with correct default state", () => {
+    const initialState = CellHeatmapFeature.getInitialState?.();
+    expect(initialState).toEqual({
+      columnHeatmap: {},
+      cachedMaxValue: null,
+      cachedMinValue: null,
+    });
+  });
+
+  it("should provide default options", () => {
+    const options = CellHeatmapFeature.getDefaultOptions?.(mockTable) || {};
+    expect(options.enableCellHeatmap).toBe(true);
+    expect(typeof options.onColumnHeatmapChange).toBe("function");
+  });
+
+  it("should add getGlobalHeatmap and toggleGlobalHeatmap methods to table", () => {
+    CellHeatmapFeature.createTable?.(mockTable);
+    expect(typeof mockTable.getGlobalHeatmap).toBe("function");
+    expect(typeof mockTable.toggleGlobalHeatmap).toBe("function");
+  });
+
+  it("should add methods to column", () => {
+    const mockColumn = { id: "test" } as Column<number>;
+    CellHeatmapFeature.createColumn?.(mockColumn, mockTable);
+    expect(typeof mockColumn.getCellHeatmapColor).toBe("function");
+    expect(typeof mockColumn.toggleColumnHeatmap).toBe("function");
+    expect(typeof mockColumn.getIsColumnHeatmapEnabled).toBe("function");
+  });
+
+  it("should calculate correct heatmap color", () => {
+    const mockColumn = { id: "value" } as Column<number>;
+    CellHeatmapFeature.createColumn?.(mockColumn, mockTable);
+    mockTable.setState((prev) => {
+      prev.columnHeatmap = { value: true };
+      return prev;
+    });
+
+    const color = mockColumn.getCellHeatmapColor?.(20);
+    expect(color).toMatch(/^hsla\(\d+(?:,\s*\d+%){2},\s*0\.6\)$/);
+  });
+
+  it("should toggle column heatmap", () => {
+    const mockColumn = { id: "value" } as Column<number>;
+    const onColumnHeatmapChange = vi.fn();
+    mockTable.options.onColumnHeatmapChange = onColumnHeatmapChange;
+
+    CellHeatmapFeature.createColumn?.(mockColumn, mockTable);
+    mockColumn.toggleColumnHeatmap?.();
+
+    expect(onColumnHeatmapChange).toHaveBeenCalled();
+  });
+
+  it("should handle global heatmap toggle", () => {
+    CellHeatmapFeature.createTable?.(mockTable);
+    const onColumnHeatmapChange = vi.fn();
+    mockTable.options.onColumnHeatmapChange = onColumnHeatmapChange;
+
+    mockTable.toggleGlobalHeatmap?.();
+    expect(onColumnHeatmapChange).toHaveBeenCalled();
+  });
+});

frontend/src/components/data-table/cell-heatmap/feature.ts

@@ -0,0 +1,202 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import {
+  type TableFeature,
+  type RowData,
+  type Table,
+  type Column,
+  type Updater,
+  makeStateUpdater,
+} from "@tanstack/react-table";
+import type { CellHeatmapTableState, CellHeatmapOptions } from "./types";
+
+export const CellHeatmapFeature: TableFeature = {
+  getInitialState: (state): CellHeatmapTableState => {
+    return {
+      columnHeatmap: {},
+      cachedMaxValue: null,
+      cachedMinValue: null,
+      ...state,
+    };
+  },
+
+  getDefaultOptions: <TData extends RowData>(
+    table: Table<TData>,
+  ): CellHeatmapOptions => {
+    return {
+      enableCellHeatmap: true,
+      onColumnHeatmapChange: makeStateUpdater("columnHeatmap", table),
+    };
+  },
+
+  createTable: <TData extends RowData>(table: Table<TData>) => {
+    table.getGlobalHeatmap = () => {
+      return Object.values(table.getState().columnHeatmap).some(Boolean);
+    };
+
+    table.toggleGlobalHeatmap = () => {
+      const allColumns = table.getAllColumns();
+      const { columnHeatmap } = table.getState();
+      const hasAnyEnabled = Object.values(columnHeatmap).some(Boolean);
+
+      if (hasAnyEnabled) {
+        // Disable all columns
+        table.options.onColumnHeatmapChange?.(
+          Object.fromEntries(allColumns.map((column) => [column.id, false])),
+        );
+      } else {
+        // Enable all columns
+        table.options.onColumnHeatmapChange?.(
+          Object.fromEntries(allColumns.map((column) => [column.id, true])),
+        );
+      }
+
+      table.setState((old) => ({
+        ...old,
+        cachedMinValue: null,
+        cachedMaxValue: null,
+      }));
+    };
+  },
+
+  createColumn: <TData extends RowData>(
+    column: Column<TData>,
+    table: Table<TData>,
+  ) => {
+    // Clear min/max cache when a column is added or removed
+    table.setState((old) => ({
+      ...old,
+      cachedMinValue: null,
+      cachedMaxValue: null,
+    }));
+
+    column.getCellHeatmapColor = (cellValue: unknown) => {
+      const state = table.getState();
+      const isColumnHeatmapEnabled = Boolean(state.columnHeatmap[column.id]);
+      if (!isColumnHeatmapEnabled || typeof cellValue !== "number") {
+        return "";
+      }
+
+      // Get all numeric values from enabled columns
+      let minValue = state.cachedMinValue;
+      let maxValue = state.cachedMaxValue;
+
+      if (minValue == null || maxValue == null) {
+        const { min, max } = getMaxMinValue(table);
+        minValue = min;
+        maxValue = max;
+        table.setState((old) => ({
+          ...old,
+          cachedMinValue: minValue,
+          cachedMaxValue: maxValue,
+        }));
+      }
+
+      const isDarkMode =
+        typeof window !== "undefined" &&
+        "matchMedia" in window &&
+        typeof window.matchMedia === "function" &&
+        window.matchMedia("(prefers-color-scheme: dark)").matches;
+
+      const colorStops = isDarkMode
+        ? [
+            { hue: 210, saturation: 100, lightness: 30 }, // Darker Blue
+            { hue: 199, saturation: 95, lightness: 33 }, // Darker Cyan
+            { hue: 172, saturation: 66, lightness: 30 }, // Darker Teal
+            { hue: 158, saturation: 64, lightness: 32 }, // Darker Green
+            { hue: 142, saturation: 71, lightness: 25 }, // Darker Lime
+            { hue: 47, saturation: 96, lightness: 33 }, // Darker Yellow
+            { hue: 21, saturation: 90, lightness: 28 }, // Darker Orange
+            { hue: 0, saturation: 84, lightness: 40 }, // Darker Red
+          ]
+        : [
+            { hue: 210, saturation: 100, lightness: 50 }, // Blue-500
+            { hue: 199, saturation: 95, lightness: 53 }, // Cyan-500
+            { hue: 172, saturation: 66, lightness: 50 }, // Teal-500
+            { hue: 158, saturation: 64, lightness: 52 }, // Green-500
+            { hue: 142, saturation: 71, lightness: 45 }, // Lime-600
+            { hue: 47, saturation: 96, lightness: 53 }, // Yellow-400
+            { hue: 21, saturation: 90, lightness: 48 }, // Orange-500
+            { hue: 0, saturation: 84, lightness: 60 }, // Red-500
+          ];
+
+      // Normalize the cellValue
+      const normalized = (cellValue - minValue) / (maxValue - minValue);
+
+      const index = Math.min(
+        Math.floor(normalized * (colorStops.length - 1)),
+        colorStops.length - 2,
+      );
+      const t = normalized * (colorStops.length - 1) - index;
+
+      const c1 = colorStops[index];
+      const c2 = colorStops[index + 1];
+
+      if (!c1 || !c2) {
+        return "";
+      }
+
+      const hue = Math.round(c1.hue + t * (c2.hue - c1.hue));
+      const saturation = Math.round(
+        c1.saturation + t * (c2.saturation - c1.saturation),
+      );
+      const lightness = Math.round(
+        c1.lightness + t * (c2.lightness - c1.lightness),
+      );
+
+      return `hsla(${hue}, ${saturation}%, ${lightness}%, 0.6)`;
+    };
+
+    column.toggleColumnHeatmap = (value?: boolean) => {
+      const safeUpdater: Updater<CellHeatmapTableState["columnHeatmap"]> = (
+        old,
+      ) => {
+        const prevValue = old[column.id];
+        if (value !== undefined) {
+          return {
+            ...old,
+            [column.id]: value,
+          };
+        }
+
+        return {
+          ...old,
+          [column.id]: !prevValue,
+        };
+      };
+
+      table.options.onColumnHeatmapChange?.(safeUpdater);
+      table.setState((old) => ({
+        ...old,
+        cachedMinValue: null,
+        cachedMaxValue: null,
+      }));
+    };
+
+    column.getIsColumnHeatmapEnabled = () => {
+      return table.getState().columnHeatmap[column.id] || false;
+    };
+  },
+};
+
+function getMaxMinValue<TData extends RowData>(table: Table<TData>) {
+  const { columnHeatmap } = table.getState();
+  const enabledColumnsSet = new Set(
+    Object.keys(columnHeatmap).filter((key) => columnHeatmap[key]),
+  );
+  const values: number[] = [];
+  for (const row of table.getRowModel().rows) {
+    for (const column of table.getAllColumns()) {
+      if (enabledColumnsSet.has(column.id)) {
+        const cellValue = row.getValue(column.id);
+        if (typeof cellValue === "number" && !Number.isNaN(cellValue)) {
+          values.push(cellValue);
+        }
+      }
+    }
+  }
+
+  return {
+    min: Math.min(...values),
+    max: Math.max(...values),
+  };
+}

frontend/src/components/data-table/cell-heatmap/types.ts

@@ -0,0 +1,40 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+/* eslint-disable @typescript-eslint/no-empty-interface */
+import type { RowData, Updater } from "@tanstack/react-table";
+
+export interface CellHeatmapTableState {
+  columnHeatmap: Record<string, boolean>;
+  cachedMinValue?: number | null;
+  cachedMaxValue?: number | null;
+}
+
+export interface CellHeatmapOptions {
+  enableCellHeatmap?: boolean;
+  onGlobalHeatmapChange?: (updater: Updater<boolean>) => void;
+  onColumnHeatmapChange?: (updater: Updater<Record<string, boolean>>) => void;
+}
+
+export interface CellHeatmapState {
+  global: boolean;
+  columns: Record<string, boolean>;
+}
+
+// Use declaration merging to add our new feature APIs
+declare module "@tanstack/react-table" {
+  interface TableState extends CellHeatmapTableState {}
+  interface InitialTableState extends CellHeatmapTableState {}
+
+  interface TableOptionsResolved<TData extends RowData>
+    extends CellHeatmapOptions {}
+
+  interface Table<TData extends RowData> {
+    getGlobalHeatmap?: () => boolean;
+    toggleGlobalHeatmap?: (value?: boolean) => void;
+  }
+
+  interface Column<TData extends RowData> {
+    getCellHeatmapColor?: (cellValue: unknown) => string;
+    toggleColumnHeatmap?: (value?: boolean) => void;
+    getIsColumnHeatmapEnabled?: () => boolean;
+  }
+}