@marimo-team/frontend 0.22.5-dev12 → 0.22.5-dev14

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

@@ -2,7 +2,13 @@
 "use no memo";
 
 import type { Column, Table } from "@tanstack/react-table";
-import { FilterIcon, MinusIcon, TextIcon, XIcon } from "lucide-react";
+import {
+  EllipsisIcon,
+  FilterIcon,
+  MinusIcon,
+  TextIcon,
+  XIcon,
+} from "lucide-react";
 import { useMemo, useRef, useState } from "react";
 import { useLocale } from "react-aria";
 import {
@@ -69,7 +75,7 @@ interface DataTableColumnHeaderProps<
 > extends React.HTMLAttributes<HTMLDivElement> {
   column: Column<TData, TValue>;
   header: React.ReactNode;
-  justify?: "left" | "center" | "right";
+  subheader?: React.ReactNode;
   calculateTopKRows?: CalculateTopKRows;
   table?: Table<TData>;
 }
@@ -77,7 +83,7 @@ interface DataTableColumnHeaderProps<
 export const DataTableColumnHeader = <TData, TValue>({
   column,
   header,
-  justify,
+  subheader,
   className,
   calculateTopKRows,
   table,
@@ -92,49 +98,51 @@ export const DataTableColumnHeader = <TData, TValue>({
 
   // No sorting or filtering
   if (!column.getCanSort() && !column.getCanFilter()) {
-    return <div className={cn(className)}>{header}</div>;
+    return (
+      <div className={cn(className)}>
+        {header}
+        {subheader}
+      </div>
+    );
   }
 
   const hasFilter = column.getFilterValue() !== undefined;
-  const hideIcon = !column.getIsSorted() && !hasFilter;
 
   return (
     <>
-      <DropdownMenu modal={false}>
-        <DropdownMenuTrigger asChild={true}>
-          <div
-            className={cn(
-              "group flex items-center my-1 space-between w-full select-none gap-2 border hover:border-border border-transparent hover:bg-(--slate-3) data-[state=open]:bg-(--slate-3) data-[state=open]:border-border rounded px-1 -mx-1",
-              justify === "right" && "flex-row-reverse",
-              className,
-            )}
-            data-testid="data-table-sort-button"
-          >
-            <span className="flex-1">{header}</span>
-            <span
-              className={cn(
-                "h-5 py-1 px-1",
-                hideIcon &&
-                  "invisible group-hover:visible data-[state=open]:visible",
-              )}
-            >
-              {renderSortFilterIcon(column)}
-            </span>
-          </div>
-        </DropdownMenuTrigger>
-        <DropdownMenuContent align="start">
-          {renderDataType(column)}
-          {renderSorts(column, table)}
-          {renderCopyColumn(column)}
-          {renderColumnPinning(column)}
-          {renderColumnWrapping(column)}
-          {renderFormatOptions(column, locale)}
-          <DropdownMenuSeparator />
-          {renderMenuItemFilter(column)}
-          {renderFilterByValues(column, setIsFilterValueOpen)}
-          {hasFilter && <ClearFilterMenuItem column={column} />}
-        </DropdownMenuContent>
-      </DropdownMenu>
+      <div
+        className={cn("group flex flex-col my-1 w-full select-none", className)}
+      >
+        <div className="flex items-center gap-1">
+          <span>{header}</span>
+          {column.getCanSort() && <SortButton column={column} />}
+          <DropdownMenu modal={false}>
+            <DropdownMenuTrigger asChild={true}>
+              <button
+                type="button"
+                className="inline-flex items-center justify-center h-5 w-5 rounded hover:bg-(--slate-4) text-muted-foreground opacity-0 group-hover:opacity-100 focus:opacity-100 group-focus-within:opacity-100 data-[state=open]:opacity-100 data-[state=open]:text-accent-foreground"
+                aria-label="Column options"
+                data-testid="data-table-column-menu-button"
+              >
+                <EllipsisIcon className="h-3.5 w-3.5" />
+              </button>
+            </DropdownMenuTrigger>
+            <DropdownMenuContent align="start">
+              {renderDataType(column)}
+              {renderSorts(column, table)}
+              {renderCopyColumn(column)}
+              {renderColumnPinning(column)}
+              {renderColumnWrapping(column)}
+              {renderFormatOptions(column, locale)}
+              <DropdownMenuSeparator />
+              {renderMenuItemFilter(column)}
+              {renderFilterByValues(column, setIsFilterValueOpen)}
+              {hasFilter && <ClearFilterMenuItem column={column} />}
+            </DropdownMenuContent>
+          </DropdownMenu>
+        </div>
+        {subheader}
+      </div>
       {isFilterValueOpen && (
         <PopoverFilterByValues
           setIsFilterValueOpen={setIsFilterValueOpen}
@@ -146,28 +154,45 @@ export const DataTableColumnHeader = <TData, TValue>({
   );
 };
 
-export const DataTableColumnHeaderWithSummary = <TData, TValue>({
+const SortButton = <TData, TValue>({
   column,
-  header,
-  summary,
-  className,
-}: DataTableColumnHeaderProps<TData, TValue> & {
-  summary: React.ReactNode;
+}: {
+  column: Column<TData, TValue>;
 }) => {
+  const sortDirection = column.getIsSorted();
+
+  const handleClick = (e: React.MouseEvent) => {
+    e.stopPropagation();
+    if (!sortDirection) {
+      column.toggleSorting(false, true); // asc
+    } else if (sortDirection === "asc") {
+      column.toggleSorting(true, true); // desc
+    } else {
+      column.clearSorting();
+    }
+  };
+
   return (
-    <div
+    <button
+      type="button"
+      onClick={handleClick}
       className={cn(
-        "flex flex-col h-full pt-0.5 pb-3 justify-between items-start",
-        className,
+        "inline-flex items-center justify-center h-5 w-5 rounded hover:bg-(--slate-4)",
+        sortDirection
+          ? "text-accent-foreground"
+          : "text-muted-foreground opacity-0 group-hover:opacity-100 focus:opacity-100 group-focus-within:opacity-100",
       )}
+      aria-label={
+        sortDirection === "asc"
+          ? "Sorted ascending, click to sort descending"
+          : sortDirection === "desc"
+            ? "Sorted descending, click to clear sort"
+            : "Sort column ascending"
+      }
+      data-testid="data-table-sort-button"
     >
-      <DataTableColumnHeader
-        column={column}
-        header={header}
-        className={className}
-      />
-      {summary}
-    </div>
+      {renderSortFilterIcon(column)}
+    </button>
   );
 };
 

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

@@ -206,41 +206,30 @@ export function generateColumns<T>({
             </div>
           ) : null;
 
-        const justify = getJustify(key);
-
-        const headerWithType = (
-          <div
+        const headerName = (
+          <span
             className={cn(
-              "flex flex-col",
-              justify === "center" && "items-center",
-              justify === "right" && "items-end",
+              "font-bold",
+              headerTitle && "underline decoration-dotted",
             )}
           >
-            <span
-              className={cn(
-                "font-bold",
-                headerTitle && "underline decoration-dotted",
-              )}
-            >
-              {key === "" ? " " : key}
-            </span>
-            {dtypeHeader}
-          </div>
+            {key === "" ? " " : key}
+          </span>
         );
 
         const headerWithTooltip = headerTitle ? (
           <Tooltip content={headerTitle} delayDuration={300}>
-            {headerWithType}
+            {headerName}
           </Tooltip>
         ) : (
-          headerWithType
+          headerName
         );
 
         const dataTableColumnHeader = (
           <DataTableColumnHeader
             header={headerWithTooltip}
+            subheader={dtypeHeader}
             column={column}
-            justify={justify}
             calculateTopKRows={calculateTopKRows}
             table={table}
           />
@@ -255,8 +244,6 @@ export function generateColumns<T>({
           <div
             className={cn(
               "flex flex-col h-full pt-0.5 pb-3 justify-between items-start",
-              justify === "center" && "items-center",
-              justify === "right" && "items-end",
             )}
           >
             {dataTableColumnHeader}
@@ -283,13 +270,13 @@ export function generateColumns<T>({
 
         const dataType = column.columnDef.meta?.dataType;
         const isNumeric = dataType === "number" || dataType === "integer";
-        const cellStyles = getCellStyleClass(
+        const cellStyles = getCellStyleClass({
           justify,
           wrapped,
           canSelectCell,
-          isCellSelected,
+          isSelected: isCellSelected,
           isNumeric,
-        );
+        });
 
         const renderedCell = renderCellValue({
           column,
@@ -448,13 +435,19 @@ function getFilterTypeForFieldType(
   }
 }
 
-function getCellStyleClass(
-  justify: "left" | "center" | "right" | undefined,
-  wrapped: boolean | undefined,
-  canSelectCell: boolean,
-  isSelected: boolean,
-  isNumeric?: boolean,
-): string {
+function getCellStyleClass({
+  justify = "left",
+  wrapped,
+  canSelectCell,
+  isSelected,
+  isNumeric = false,
+}: {
+  justify: "left" | "center" | "right" | undefined;
+  wrapped: boolean | undefined;
+  canSelectCell: boolean;
+  isSelected: boolean;
+  isNumeric?: boolean;
+}): string {
   return cn(
     canSelectCell && "cursor-pointer",
     isSelected &&

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

@@ -47,6 +47,7 @@ import { DataTableBody, renderTableHeader } from "./renderers";
 import { TableBottomBar } from "./TableBottomBar";
 import { TableTopBar } from "./TableTopBar";
 import {
+  AUTO_WIDTH_MAX_COLUMNS,
   type DataTableSelection,
   MIN_ROWS_TO_VIRTUALIZE,
   type TooManyRows,
@@ -300,7 +301,13 @@ const DataTableInternal = <TData,>({
             isAnyPanelOpen={isAnyPanelOpen}
             downloadAs={downloadAs}
           />
-          <Table className="relative" ref={tableRef}>
+          <Table
+            className={cn(
+              "relative",
+              columns.length <= AUTO_WIDTH_MAX_COLUMNS ? "w-auto" : "w-full",
+            )}
+            ref={tableRef}
+          >
             {showLoadingBar && (
               <thead className="absolute top-0 left-0 h-[3px] w-1/2 bg-primary animate-slide" />
             )}

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

@@ -27,7 +27,7 @@ import { DataTableContextMenu } from "./context-menu";
 import { CellRangeSelectionIndicator } from "./range-focus/cell-selection-indicator";
 import { useCellRangeSelection } from "./range-focus/use-cell-range-selection";
 import { useScrollIntoViewOnFocus } from "./range-focus/use-scroll-into-view";
-import { TABLE_ROW_HEIGHT_PX } from "./types";
+import { AUTO_WIDTH_MAX_COLUMNS, TABLE_ROW_HEIGHT_PX } from "./types";
 import { stringifyUnknownValue } from "./utils";
 
 export function renderTableHeader<TData>(
@@ -46,7 +46,7 @@ export function renderTableHeader<TData>(
           <TableHead
             key={header.id}
             className={cn(
-              "h-auto min-h-10 whitespace-pre align-top",
+              "h-auto min-h-10 whitespace-pre align-top border-r border-r-border/75",
               className,
             )}
             style={style}
@@ -69,6 +69,13 @@ export function renderTableHeader<TData>(
         {renderHeaderGroup(table.getLeftHeaderGroups())}
         {renderHeaderGroup(table.getCenterHeaderGroups())}
         {renderHeaderGroup(table.getRightHeaderGroups())}
+        {table.getAllColumns().length <= AUTO_WIDTH_MAX_COLUMNS && (
+          <th
+            className="w-full border-0"
+            aria-hidden="true"
+            role="presentation"
+          />
+        )}
       </TableRow>
     </TableHeader>
   );
@@ -163,7 +170,7 @@ export const DataTableBody = <TData,>({
           {...getCellDomProps(cell.id)}
           key={cell.id}
           className={cn(
-            "whitespace-pre truncate max-w-[300px] outline-hidden",
+            "whitespace-pre truncate max-w-[300px] outline-hidden border-r border-r-border/75",
             cell.column.getColumnWrapping &&
               cell.column.getColumnWrapping?.() === "wrap" &&
               COLUMN_WRAPPING_STYLES,
@@ -230,15 +237,21 @@ export const DataTableBody = <TData,>({
         {renderCells(row.getLeftVisibleCells())}
         {renderCells(row.getCenterVisibleCells())}
         {renderCells(row.getRightVisibleCells())}
+        {columns.length <= AUTO_WIDTH_MAX_COLUMNS && (
+          <td className="border-0" aria-hidden="true" role="presentation" />
+        )}
       </TableRow>
     );
   };
 
+  const hasFillerColumn = columns.length <= AUTO_WIDTH_MAX_COLUMNS;
+  const totalColSpan = columns.length + (hasFillerColumn ? 1 : 0);
+
   const renderRows = () => {
     if (rows.length === 0) {
       return (
         <TableRow>
-          <TableCell colSpan={columns.length} className="h-24 text-center">
+          <TableCell colSpan={totalColSpan} className="h-24 text-center">
             No results.
           </TableCell>
         </TableRow>
@@ -255,7 +268,7 @@ export const DataTableBody = <TData,>({
               data-virtual-spacer=""
               style={{ height: virtualItems[0].start }}
             >
-              <td colSpan={columns.length} />
+              <td colSpan={totalColSpan} />
             </tr>
           )}
           {virtualItems.map((vItem) => renderRow(rows[vItem.index]))}
@@ -266,7 +279,7 @@ export const DataTableBody = <TData,>({
                 height: totalSize - (virtualItems.at(-1)?.end ?? totalSize),
               }}
             >
-              <td colSpan={columns.length} />
+              <td colSpan={totalColSpan} />
             </tr>
           )}
         </>

package/src/components/data-table/types.ts

@@ -16,6 +16,10 @@ declare module "@tanstack/react-table" {
 export const TABLE_ROW_HEIGHT_PX = 24;
 export const TABLE_HEADER_HEIGHT_PX = 40;
 
+// Below this column count, the table uses w-auto with a filler column
+// to prevent columns from stretching unnecessarily
+export const AUTO_WIDTH_MAX_COLUMNS = 4;
+
 // Default number of visible rows when virtualizing without an explicit maxHeight.
 export const DEFAULT_VIRTUAL_ROWS = 15;
 

package/src/components/editor/actions/useNotebookActions.tsx

@@ -17,7 +17,6 @@ import {
   EyeOffIcon,
   FastForwardIcon,
   FileIcon,
-  FilePlus2Icon,
   Files,
   FileTextIcon,
   FolderDownIcon,
@@ -44,6 +43,7 @@ import {
 } from "lucide-react";
 import { settingDialogAtom } from "@/components/app-config/state";
 import { MarkdownIcon } from "@/components/editor/cell/code/icons";
+import { MarimoPlusIcon } from "@/components/icons/marimo-icons";
 import { useImperativeModal } from "@/components/modal/ImperativeModal";
 import { renderShortcut } from "@/components/shortcuts/renderShortcut";
 import { PairWithAgentModal } from "@/components/editor/actions/pair-with-agent-modal";
@@ -640,7 +640,7 @@ export function useNotebookActions() {
     },
 
     {
-      icon: <FilePlus2Icon size={14} strokeWidth={1.5} />,
+      icon: <MarimoPlusIcon size={14} strokeWidth={1.5} />,
       label: "New notebook",
       // If file is in the url, then we ran `marimo edit`
       // without a specific file

package/src/components/editor/errors/traceback-modal.tsx

@@ -12,6 +12,8 @@ import {
 } from "@/components/ui/alert-dialog";
 import { CopyIcon } from "lucide-react";
 import { toast } from "@/components/ui/use-toast";
+import { renderHTML } from "@/plugins/core/RenderHTML";
+import { sanitizeHtml } from "@/plugins/core/sanitize-html";
 
 interface TracebackModalProps {
   isOpen: boolean;
@@ -29,7 +31,7 @@ export const TracebackModal: React.FC<TracebackModalProps> = ({
   const handleCopy = async () => {
     // Strip HTML tags for clipboard
     const tempDiv = document.createElement("div");
-    tempDiv.innerHTML = traceback;
+    tempDiv.innerHTML = sanitizeHtml(traceback);
     const textContent = tempDiv.textContent || tempDiv.innerText || "";
 
     try {
@@ -73,10 +75,9 @@ export const TracebackModal: React.FC<TracebackModalProps> = ({
               Copy
             </Button>
           </div>
-          <div
-            className="font-code text-sm p-4 bg-muted rounded border overflow-auto max-h-[50vh] cursor-text select-text"
-            dangerouslySetInnerHTML={{ __html: traceback }}
-          />
+          <div className="font-code text-sm p-4 bg-muted rounded border overflow-auto max-h-[50vh] cursor-text select-text">
+            {renderHTML({ html: traceback })}
+          </div>
         </div>
         <AlertDialogFooter>
           <AlertDialogAction onClick={onClose}>Close</AlertDialogAction>

package/src/components/editor/output/MarimoErrorOutput.tsx

@@ -25,6 +25,7 @@ import { useChromeActions } from "../chrome/state";
 import { AutoFixButton } from "../errors/auto-fix";
 import { CellLinkError } from "../links/cell-link";
 import { processTextForUrls } from "./console/text-rendering";
+import { renderHTML } from "@/plugins/core/RenderHTML";
 
 const Tip = (props: {
   title?: string;
@@ -486,13 +487,9 @@ export const MarimoErrorOutput = ({
                       {processTextForUrls(error.msg, `exception-${idx}`)}
                     </p>
                     {"traceback" in error && error.traceback ? (
-                      <div
-                        className="font-code text-sm mt-2 p-3 bg-muted rounded border overflow-auto max-h-[50vh] cursor-text select-text"
-                        // biome-ignore lint/security/noDangerouslySetInnerHtml: traceback from backend
-                        dangerouslySetInnerHTML={{
-                          __html: error.traceback,
-                        }}
-                      />
+                      <div className="font-code text-sm mt-2 p-3 bg-muted rounded border overflow-auto max-h-[50vh] cursor-text select-text">
+                        {renderHTML({ html: error.traceback })}
+                      </div>
                     ) : (
                       <div className="text-muted-foreground mt-2">
                         See the console area for a traceback.
@@ -504,13 +501,9 @@ export const MarimoErrorOutput = ({
                     {processTextForUrls(error.msg, `exception-${idx}`)}
                     <CellLinkError cellId={error.raising_cell} />
                     {"traceback" in error && error.traceback && (
-                      <div
-                        className="font-code text-sm mt-2 p-3 bg-muted rounded border overflow-auto max-h-[50vh] cursor-text select-text"
-                        // biome-ignore lint/security/noDangerouslySetInnerHtml: traceback from backend
-                        dangerouslySetInnerHTML={{
-                          __html: error.traceback,
-                        }}
-                      />
+                      <div className="font-code text-sm mt-2 p-3 bg-muted rounded border overflow-auto max-h-[50vh] cursor-text select-text">
+                        {renderHTML({ html: error.traceback })}
+                      </div>
                     )}
                   </div>
                 )}

package/src/components/editor/output/MarimoTracebackOutput.tsx

@@ -34,6 +34,7 @@ import { getRequestClient } from "@/core/network/requests";
 import { isStaticNotebook } from "@/core/static/static-state";
 import { isWasm } from "@/core/wasm/utils";
 import { renderHTML } from "@/plugins/core/RenderHTML";
+import { sanitizeHtml } from "@/plugins/core/sanitize-html";
 import { copyToClipboard } from "@/utils/copy";
 import {
   elementContainsMarimoCellFile,
@@ -173,9 +174,9 @@ export const MarimoTracebackOutput = ({
               </DropdownMenuItem>
               <DropdownMenuItem
                 onClick={() => {
-                  // Strip HTML from the traceback
+                  // Strip HTML from the traceback (sanitize first to prevent XSS)
                   const div = document.createElement("div");
-                  div.innerHTML = traceback;
+                  div.innerHTML = sanitizeHtml(traceback);
                   const textContent = div.textContent || "";
                   copyToClipboard(textContent);
                 }}
@@ -193,7 +194,7 @@ export const MarimoTracebackOutput = ({
 
 function lastLine(text: string): string {
   const el = document.createElement("div");
-  el.innerHTML = text;
+  el.innerHTML = sanitizeHtml(text);
   const lines = el.textContent?.split("\n").filter(Boolean);
   return lines?.at(-1) || "";
 }

package/src/plugins/impl/plotly/__tests__/selection.test.ts

@@ -117,6 +117,14 @@ describe("shouldHandleClickSelection", () => {
     expect(shouldHandleClickSelection([linePoint])).toBe(true);
   });
 
+  it("accepts waterfall clicks", () => {
+    const waterfallPoint = createPlotDatum({
+      data: { type: "waterfall" },
+    });
+
+    expect(shouldHandleClickSelection([waterfallPoint])).toBe(true);
+  });
+
   it("rejects non-line scatter marker clicks", () => {
     const markerPoint = createPlotDatum({
       data: { type: "scatter", mode: "markers" },
@@ -196,4 +204,18 @@ describe("extractPoints", () => {
 
     expect(extractPoints([point])).toEqual([{ x: 1, y: 2, z: 3 }]);
   });
+
+  it("returns x/y/pointIndex for waterfall clicks", () => {
+    const point = createPlotDatum({
+      x: "Revenue",
+      y: 400,
+      pointIndex: 1,
+      curveNumber: 0,
+      data: { type: "waterfall" },
+    });
+
+    expect(extractPoints([point])).toEqual([
+      { x: "Revenue", y: 400, pointIndex: 1, curveNumber: 0 },
+    ]);
+  });
 });

package/src/plugins/impl/plotly/selection.ts

@@ -227,6 +227,7 @@ export function shouldHandleClickSelection(
       type === "bar" ||
       type === "heatmap" ||
       type === "histogram" ||
+      type === "waterfall" ||
       isLinePoint(point)
     );
   });