@marimo-team/frontend 0.25.3-dev12 → 0.25.3-dev121

package/src/components/data-table/charts/charts.tsx

@@ -13,8 +13,8 @@ import {
   XIcon,
 } from "lucide-react";
 import type { JSX } from "react";
-import React, { useMemo, useRef, useState } from "react";
-import { type UseFormReturn, useForm } from "react-hook-form";
+import React, { useMemo, useState } from "react";
+import { type UseFormReturn, useForm, useWatch } from "react-hook-form";
 import useResizeObserver from "use-resize-observer";
 import { PythonIcon } from "@/components/editor/cell/code/icons";
 import { Alert, AlertDescription } from "@/components/ui/alert";
@@ -48,7 +48,12 @@ import { HeatmapForm } from "./forms/heatmap";
 import { PieForm } from "./forms/pie";
 import { LazyChart } from "./lazy-chart";
 import { ChartSchema, type ChartSchemaType, getChartDefaults } from "./schemas";
-import { getChartTabName, type TabName, tabsStorageAtom } from "./storage";
+import {
+  createTabId,
+  getChartTabName,
+  type TabName,
+  tabsStorageAtom,
+} from "./storage";
 import { ChartType } from "./types";
 
 const NEW_CHART_TYPE = "bar" as ChartType;
@@ -103,26 +108,28 @@ export const TablePanel: React.FC<TablePanelProps> = ({
 
   const [selectedTab, setSelectedTab] = useState(DEFAULT_TAB_NAME);
   const [tabCounter, setTabCounter] = useState(tabs.length);
-  const prevDisplayHeader = useRef(displayHeader);
+  const [prevDisplayHeader, setPrevDisplayHeader] = useState(displayHeader);
+
+  if (displayHeader !== prevDisplayHeader) {
+    setPrevDisplayHeader(displayHeader);
+  }
 
   // Auto-create a default chart tab when chart builder opens with no tabs
-  if (
-    displayHeader &&
-    !prevDisplayHeader.current &&
-    tabs.length === 0 &&
-    cellId
-  ) {
-    prevDisplayHeader.current = displayHeader;
+  if (displayHeader && !prevDisplayHeader && tabs.length === 0 && cellId) {
     const tabName = getChartTabName(0, NEW_CHART_TYPE);
     const newTabs = new Map(tabsMap);
     newTabs.set(cellId, [
-      { tabName, chartType: NEW_CHART_TYPE, config: getChartDefaults() },
+      {
+        id: createTabId(),
+        tabName,
+        chartType: NEW_CHART_TYPE,
+        config: getChartDefaults(),
+      },
     ]);
     saveTabsMap(newTabs);
     setTabCounter(1);
     setSelectedTab(tabName);
   }
-  prevDisplayHeader.current = displayHeader;
 
   if (!displayHeader || (tabs.length === 0 && !displayHeader)) {
     return dataTable;
@@ -138,6 +145,7 @@ export const TablePanel: React.FC<TablePanelProps> = ({
     newTabs.set(cellId, [
       ...tabs,
       {
+        id: createTabId(),
         tabName,
         chartType: NEW_CHART_TYPE,
         config: getChartDefaults(),
@@ -236,9 +244,9 @@ export const TablePanel: React.FC<TablePanelProps> = ({
         >
           Table
         </TabsTrigger>
-        {tabs.map((tab, idx) => (
+        {tabs.map((tab) => (
           <TabsTrigger
-            key={idx}
+            key={tab.id}
             className="text-xs"
             value={tab.tabName}
             onClick={() => setSelectedTab(tab.tabName)}
@@ -266,7 +274,7 @@ export const TablePanel: React.FC<TablePanelProps> = ({
       <TabsContent className="mt-1 overflow-hidden" value={DEFAULT_TAB_NAME}>
         {dataTable}
       </TabsContent>
-      {tabs.map((tab, idx) => {
+      {tabs.map((tab) => {
         const saveChart = (formValues: ChartSchemaType) => {
           saveTabChart({
             tabName: tab.tabName,
@@ -278,7 +286,11 @@ export const TablePanel: React.FC<TablePanelProps> = ({
           saveTabChartType(tab.tabName, chartType);
         };
         return (
-          <TabsContent key={idx} value={tab.tabName} className="h-[400px] mt-0">
+          <TabsContent
+            key={tab.id}
+            value={tab.tabName}
+            className="h-[400px] mt-0"
+          >
             <ChartPanel
               tableData={data}
               chartConfig={tab.config}
@@ -320,9 +332,10 @@ export const ChartPanel: React.FC<{
   saveChartType,
   getDataUrl,
   fieldTypes,
-  hasSchema = Boolean(fieldTypes?.length),
+  hasSchema: providedHasSchema,
   isLargeDataset,
 }) => {
+  const hasSchema = providedHasSchema ?? Boolean(fieldTypes?.length);
   const { theme } = useTheme();
   const form = useForm<ChartSchemaType>({
     defaultValues: chartConfig ?? getChartDefaults(),
@@ -370,20 +383,17 @@ export const ChartPanel: React.FC<{
     // Re-run when the data table changes
   }, [tableData, renderLargeCharts, csvFieldTypes]);
 
-  const formValues = form.watch();
+  // useWatch returns a new object only when a value changes. It is typed as
+  // deeply partial, but the form always starts from complete defaults.
+  const formValues = useWatch({ control: form.control }) as ChartSchemaType;
 
-  // This ensures the chart re-renders when the actual values change
-  const memoizedFormValues = useMemo(() => {
-    return structuredClone(formValues);
-  }, [formValues]);
-
-  const specWithoutData = createSpecWithoutData(
-    selectedChartType,
-    memoizedFormValues,
+  const specWithoutData = createSpecWithoutData({
+    chartType: selectedChartType,
+    formValues,
     theme,
-    "container",
-    CHART_HEIGHT,
-  );
+    width: "container",
+    height: CHART_HEIGHT,
+  });
 
   // Prevent unnecessary re-renders of the chart
   const memoizedChart = useMemo(() => {

package/src/components/data-table/charts/storage.ts

@@ -7,13 +7,16 @@ import { Logger } from "@/utils/Logger";
 import { NotebookScopedLocalStorage } from "@/utils/storage/typed";
 import { capitalize } from "@/utils/strings";
 import type { TypedString } from "@/utils/typed";
+import { generateUUID } from "@/utils/uuid";
 import { ChartSchema, type ChartSchemaType } from "./schemas";
 import type { ChartType } from "./types";
 
 export type TabName = TypedString<"TabName">;
+export type TabId = TypedString<"TabId">;
 export const KEY = "marimo:charts:v2";
 
 interface TabStorage {
+  id: TabId;
   tabName: TabName; // unique within cell
   chartType: ChartType;
   config: ChartSchemaType;
@@ -21,11 +24,15 @@ interface TabStorage {
 
 type TabStorageMap = Map<CellId, TabStorage[]>;
 
-const TabStorageSchema = z.object({
-  tabName: z.string().transform((name) => name as TabName),
-  chartType: z.string().transform((type) => type as ChartType),
-  config: ChartSchema,
-});
+const TabStorageSchema = z
+  .object({
+    id: z.string().optional(),
+    tabName: z.string().transform((name) => name as TabName),
+    chartType: z.string().transform((type) => type as ChartType),
+    config: ChartSchema,
+  })
+  // Stored tabs without an id use their name, which is unique within a cell.
+  .transform((tab) => ({ ...tab, id: (tab.id ?? tab.tabName) as TabId }));
 const TabStorageEntriesSchema = z.array(
   z.tuple([
     z.string().transform((name) => name as CellId),
@@ -64,6 +71,10 @@ export const tabsStorageAtom = atomWithStorage<TabStorageMap>(
   mapStorage,
 );
 
+export function createTabId(): TabId {
+  return generateUUID() as TabId;
+}
+
 /**
  * Convenience function to get the tab name for a given tab number and chart type
  */

package/src/components/data-table/column-explorer-panel/column-explorer.tsx

@@ -345,7 +345,7 @@ const ColumnPreview = ({
   const chart = chart_spec && renderChart(chart_spec, theme);
 
   const addDataframeChart = chart_code && (
-    <AddDataframeChart chartCode={chart_code} />
+    <AddDataframeChart chartCode={chart_code} columnName={columnName} />
   );
 
   return (

package/src/components/data-table/column-summary/chart-spec-model.tsx

@@ -119,16 +119,16 @@ export class ColumnChartSpecModel<T> {
   private columnBinValues = new Map<ColumnName, BinValues>();
   private columnValueCounts = new Map<ColumnName, ValueCounts>();
 
-  public static readonly EMPTY = new ColumnChartSpecModel(
-    [],
-    new Map(),
-    {},
-    {},
-    {},
-    {
+  public static readonly EMPTY = new ColumnChartSpecModel({
+    data: [],
+    fieldTypes: new Map(),
+    stats: {},
+    binValues: {},
+    valueCounts: {},
+    opts: {
       includeCharts: false,
     },
-  );
+  });
 
   private dataSpec: TopLevelSpec["data"];
 
@@ -144,16 +144,23 @@ export class ColumnChartSpecModel<T> {
     includeCharts: boolean;
   };
 
-  public constructor(
-    data: T[] | string,
-    fieldTypes: FieldTypes,
-    stats: Record<ColumnName, Partial<ColumnHeaderStats>>,
-    binValues: Record<ColumnName, BinValues>,
-    valueCounts: Record<ColumnName, ValueCounts>,
+  public constructor({
+    data,
+    fieldTypes,
+    stats,
+    binValues,
+    valueCounts,
+    opts,
+  }: {
+    data: T[] | string;
+    fieldTypes: FieldTypes;
+    stats: Record<ColumnName, Partial<ColumnHeaderStats>>;
+    binValues: Record<ColumnName, BinValues>;
+    valueCounts: Record<ColumnName, ValueCounts>;
     opts: {
       includeCharts: boolean;
-    },
-  ) {
+    };
+  }) {
     this.fieldTypes = fieldTypes;
     this.stats = stats;
     this.binValues = binValues;

package/src/components/data-table/column-summary/column-summary.tsx

@@ -16,7 +16,7 @@ export const ColumnChartContext = React.createContext<
   ColumnChartSpecModel<unknown>
 >(ColumnChartSpecModel.EMPTY);
 
-interface Props<TData, TValue> {
+interface Props {
   columnId: string;
 }
 
@@ -24,9 +24,7 @@ interface Props<TData, TValue> {
 // for all calls with the same key.
 const batchedLoader = createBatchedLoader();
 
-export const TableColumnSummary = <TData, TValue>({
-  columnId,
-}: Props<TData, TValue>) => {
+export const TableColumnSummary = ({ columnId }: Props) => {
   const { locale } = useLocale();
 
   const chartSpecModel = React.use(ColumnChartContext);

package/src/components/data-table/column-visibility-dropdown.tsx

@@ -6,7 +6,7 @@
 
 import type { Table } from "@tanstack/react-table";
 import { Columns3Icon, EyeIcon, EyeOffIcon, ScanEyeIcon } from "lucide-react";
-import React, { useMemo } from "react";
+import React from "react";
 import { ColumnName } from "@/components/datasources/components";
 import { Button } from "@/components/ui/button";
 import {
@@ -24,6 +24,7 @@ import {
 } from "@/components/ui/popover";
 import { type BulkAction, useSelectList } from "@/components/ui/select-core";
 import type { DataType } from "@/core/kernel/messages";
+import { Arrays } from "@/utils/arrays";
 import { cn } from "@/utils/cn";
 import { Events } from "@/utils/events";
 import { smartMatchFilter } from "@/utils/smartMatch";
@@ -67,9 +68,8 @@ export const ColumnVisibilityDropdown = <TData,>({
   const hideableIds = userColumns
     .filter((column) => column.getCanHide())
     .map((column) => column.id);
-  const visibleHideableCount = hideableIds.filter(
-    (id) => !hiddenIds.includes(id),
-  ).length;
+  const visibleHideableColumnIds = Arrays.difference(hideableIds, hiddenIds);
+  const visibleHideableCount = visibleHideableColumnIds.length;
 
   const applyHidden = (next: string[] | string | null) => {
     const hidden = new Set(Array.isArray(next) ? next : []);
@@ -106,15 +106,10 @@ export const ColumnVisibilityDropdown = <TData,>({
           .filter((option) => !option.disabled)
           .map((option) => option.value)
       : [];
-  const showOnlyColumnState = useMemo(
-    () => ({
-      visibleHideableColumnIds: hideableIds.filter(
-        (id) => !hiddenIds.includes(id),
-      ),
-      hideableColumnIdSet: new Set(hideableIds),
-    }),
-    [hideableIds, hiddenIds],
-  );
+  const showOnlyColumnState = {
+    visibleHideableColumnIds,
+    hideableColumnIdSet: new Set(hideableIds),
+  };
   const hasSearchBulkActions =
     showOnlyMatchIds.length > 0 || matchingActions.length > 0;
 

package/src/components/data-table/columns.tsx

@@ -41,6 +41,7 @@ import {
   isNumericType,
 } from "./types";
 import { SentinelCell, WhitespaceMarkers } from "./sentinel-cell";
+import { highlightMatches } from "./search-highlight";
 import { detectSentinel, splitLeadingTrailingWhitespace } from "./utils";
 import { uniformSample } from "./uniformSample";
 import { MarkdownUrlDetector, UrlDetector } from "./url-detector";
@@ -135,6 +136,7 @@ export function generateColumns<T>({
 }): ColumnDef<T>[] {
   // Row-headers are typically index columns
   const rowHeadersSet = new Set(rowHeaders.map(([columnName]) => columnName));
+  const wrappedColumnsSet = new Set(wrappedColumns);
 
   const typesByColumn = Maps.keyBy(fieldTypes, (entry) => entry[0]);
 
@@ -277,7 +279,7 @@ export function generateColumns<T>({
       );
     },
 
-    cell: ({ column, renderValue, getValue, cell }) => {
+    cell: ({ column, renderValue, getValue, cell, table }) => {
       function selectCell() {
         if (selection !== "single-cell" && selection !== "multi-cell") {
           return;
@@ -287,7 +289,7 @@ export function generateColumns<T>({
       }
 
       const justify = getJustify(key);
-      const wrapped = wrappedColumns?.includes(key);
+      const wrapped = wrappedColumnsSet.has(key);
       const isCellSelected = cell?.getIsSelected?.() || false;
       const canSelectCell =
         (selection === "single-cell" || selection === "multi-cell") &&
@@ -309,6 +311,7 @@ export function generateColumns<T>({
         getValue,
         selectCell,
         cellStyles,
+        highlightPattern: table.options.meta?.highlightPattern,
       });
 
       // Row headers are bold
@@ -585,12 +588,14 @@ function renderDate({
   dtype,
   displayValue,
   locale,
+  highlightPattern,
 }: {
   value: Date;
   dataType?: DataType;
   dtype?: string;
   displayValue?: string;
   locale: string;
+  highlightPattern?: RegExp;
 }): React.ReactNode {
   const type = dataType === "date" ? "date" : "datetime";
   const timezone = extractTimezone(dtype);
@@ -599,7 +604,7 @@ function renderDate({
 
   return (
     <DatePopover date={value} type={type}>
-      {exactValue}
+      {highlightMatches(exactValue, highlightPattern)}
     </DatePopover>
   );
 }
@@ -610,12 +615,14 @@ export function renderCellValue<TData, TValue>({
   getValue,
   selectCell,
   cellStyles,
+  highlightPattern,
 }: {
   column: Column<TData, TValue>;
   renderValue: () => TValue | null;
   getValue: () => TValue;
   selectCell?: () => void;
   cellStyles?: string;
+  highlightPattern?: RegExp;
 }) {
   const value = getValue();
   const format = column.getColumnFormatting?.();
@@ -625,6 +632,28 @@ export function renderCellValue<TData, TValue>({
 
   const isWrapped = column.getColumnWrapping?.() === "wrap";
 
+  const textCell = (text: string) => (
+    <div onClick={selectCell} className={cellStyles}>
+      {highlightMatches(text, highlightPattern)}
+    </div>
+  );
+  const dateCell = (date: Date, displayValue?: string) => (
+    <div onClick={selectCell} className={cellStyles}>
+      <WithLocale>
+        {(locale) =>
+          renderDate({
+            value: date,
+            dataType,
+            dtype,
+            displayValue,
+            locale,
+            highlightPattern,
+          })
+        }
+      </WithLocale>
+    </div>
+  );
+
   // Sentinel values (null, whitespace, NaN, Infinity, NaT) rendered specially.
   // Empty strings are left as-is
   const sentinel = detectSentinel(value, dataType, dtype);
@@ -640,23 +669,10 @@ export function renderCellValue<TData, TValue>({
     try {
       const date = new Date(value);
       if (!format || !isValid(date)) {
-        return (
-          <div onClick={selectCell} className={cellStyles}>
-            {value}
-          </div>
-        );
+        return textCell(value);
       }
 
-      const displayValue = String(column.applyColumnFormatting(value));
-      return (
-        <div onClick={selectCell} className={cellStyles}>
-          <WithLocale>
-            {(locale) =>
-              renderDate({ value: date, dataType, dtype, displayValue, locale })
-            }
-          </WithLocale>
-        </div>
-      );
+      return dateCell(date, String(column.applyColumnFormatting(value)));
     } catch (error) {
       Logger.error("Error parsing datetime, fallback to string", error);
     }
@@ -672,17 +688,9 @@ export function renderCellValue<TData, TValue>({
     }
 
     // e.g. 2010-10-07 17:15:00
-    const displayValue = format
-      ? String(column.applyColumnFormatting(value))
-      : undefined;
-    return (
-      <div onClick={selectCell} className={cellStyles}>
-        <WithLocale>
-          {(locale) =>
-            renderDate({ value, dataType, dtype, displayValue, locale })
-          }
-        </WithLocale>
-      </div>
+    return dateCell(
+      value,
+      format ? String(column.applyColumnFormatting(value)) : undefined,
     );
   }
 
@@ -728,7 +736,7 @@ export function renderCellValue<TData, TValue>({
           className={cn(cellStyles, isWrapped && COLUMN_WRAPPING_STYLES)}
         >
           <WhitespaceMarkers value={leading} />
-          <UrlDetector parts={parts} />
+          <UrlDetector parts={parts} highlightPattern={highlightPattern} />
           <WhitespaceMarkers value={trailing} />
         </div>
       );
@@ -754,11 +762,7 @@ export function renderCellValue<TData, TValue>({
   }
 
   if (format) {
-    return (
-      <div onClick={selectCell} className={cellStyles}>
-        {column.applyColumnFormatting(value)}
-      </div>
-    );
+    return textCell(String(column.applyColumnFormatting(value)));
   }
 
   // Format to the correct locale
@@ -768,26 +772,19 @@ export function renderCellValue<TData, TValue>({
         <LocaleNumber
           value={value}
           minFractionDigits={column.columnDef.meta?.minFractionDigits}
+          highlightPattern={highlightPattern}
         />
       </div>
     );
   }
 
   if (typeof value === "boolean") {
-    return (
-      <div onClick={selectCell} className={cellStyles}>
-        {value ? "True" : "False"}
-      </div>
-    );
+    return textCell(value ? "True" : "False");
   }
 
   if (isPrimitiveOrNullish(value)) {
     const rendered = renderValue();
-    return (
-      <div onClick={selectCell} className={cellStyles}>
-        {rendered == null ? "" : String(rendered)}
-      </div>
-    );
+    return textCell(rendered == null ? "" : String(rendered));
   }
 
   const mimeValues = getMimeValues(value);
@@ -825,14 +822,16 @@ export function renderCellValue<TData, TValue>({
 export const LocaleNumber = ({
   value,
   minFractionDigits,
+  highlightPattern,
 }: {
   value: number;
   minFractionDigits?: number;
+  highlightPattern?: RegExp;
 }) => {
   const { locale } = useLocale();
   const format = useNumberFormatter({
     minimumFractionDigits: minFractionDigits,
     maximumFractionDigits: maxFractionalDigits(locale),
   });
-  return format.format(value);
+  return highlightMatches(format.format(value), highlightPattern);
 };

package/src/components/data-table/context-menu.tsx

@@ -12,7 +12,7 @@ import { Logger } from "@/utils/Logger";
 import type { Cell } from "@tanstack/react-table";
 import { useAtomValue } from "jotai";
 import { CopyIcon, FilterIcon, SquareStack } from "lucide-react";
-import type { RefObject } from "react";
+import { type RefObject, useRef, useState } from "react";
 import useEvent from "react-use-event-hook";
 import {
   ContextMenu,
@@ -123,8 +123,21 @@ export const DataTableContextMenu = <TData,>({
   tableRef: RefObject<HTMLTableSectionElement | null>;
   copyAllCells: () => void;
 }) => {
+  const [activeCell, setActiveCell] = useState<Cell<TData, unknown> | null>(
+    null,
+  );
+  const highlightedCell = useRef<Element | null>(null);
   const handleContextMenuChange = useEvent((open: boolean) => {
+    highlightedCell.current?.classList.remove("bg-(--green-4)");
+    highlightedCell.current = null;
+    if (!open) {
+      // Forget the cell so a later right-click outside any cell can't reuse it
+      contextMenuRef.current = null;
+      return;
+    }
+
     const cell = contextMenuRef.current;
+    setActiveCell(cell);
     if (!cell) {
       return;
     }
@@ -138,42 +151,45 @@ export const DataTableContextMenu = <TData,>({
       return;
     }
 
-    if (open) {
-      cellElement.classList.add("bg-(--green-4)");
-    } else {
-      cellElement.classList.remove("bg-(--green-4)");
-    }
+    cellElement.classList.add("bg-(--green-4)");
+    highlightedCell.current = cellElement;
   });
 
   return (
-    <ContextMenu onOpenChange={handleContextMenuChange}>
-      <ContextMenuTrigger asChild={true}>{tableBody}</ContextMenuTrigger>
+    // Open via the trigger so a right-click while already open retargets the
+    // menu; Radix only reports open changes, so use it for closing.
+    <ContextMenu
+      onOpenChange={(open) => {
+        if (!open) {
+          handleContextMenuChange(false);
+        }
+      }}
+    >
+      <ContextMenuTrigger
+        asChild={true}
+        onContextMenu={() => handleContextMenuChange(true)}
+      >
+        {tableBody}
+      </ContextMenuTrigger>
       <ContextMenuPortal>
-        <CellContextMenu
-          cellRef={contextMenuRef}
-          copySelectedCells={copyAllCells}
-        />
+        {activeCell && (
+          <CellContextMenu cell={activeCell} copySelectedCells={copyAllCells} />
+        )}
       </ContextMenuPortal>
     </ContextMenu>
   );
 };
 
 export const CellContextMenu = <TData,>({
-  cellRef,
+  cell,
   copySelectedCells,
 }: {
-  cellRef: RefObject<Cell<TData, unknown> | null>;
+  cell: Cell<TData, unknown>;
   copySelectedCells: () => void;
 }) => {
   const selectedCells = useAtomValue(selectedCellsAtom);
   const multipleSelectedCells = selectedCells.size > 1;
 
-  const cell = cellRef.current;
-  if (!cell) {
-    Logger.error("No cell found in context menu");
-    return;
-  }
-
   const table = cell.getContext().table;
   const displayedValue = cell.getValue();
   const rawValue =

package/src/components/data-table/data-table.tsx

@@ -58,6 +58,7 @@ import { useColumnPinning } from "./hooks/use-column-pinning";
 import { useScrollContainerHeight } from "./hooks/use-scroll-container-height";
 import { CellSelectionProvider } from "./range-focus/provider";
 import { DataTableBody, renderTableHeader } from "./renderers";
+import { buildHighlightPattern } from "./search-highlight";
 import { TableBottomBar } from "./TableBottomBar";
 import { TableTopBar } from "./TableTopBar";
 import {
@@ -231,6 +232,11 @@ const DataTableInternal = <TData,>({
     return idx + offset;
   }
 
+  const highlightPattern = React.useMemo(
+    () => buildHighlightPattern(searchQuery),
+    [searchQuery],
+  );
+
   const table = useReactTable<TData>({
     _features: [
       ColumnPinning,
@@ -245,7 +251,7 @@ const DataTableInternal = <TData,>({
     ],
     data,
     columns,
-    meta: { rawData },
+    meta: { rawData, highlightPattern },
     getCoreRowModel: getCoreRowModel(),
     // pagination
     rowCount: totalRows === "too_many" ? undefined : totalRows,