@marimo-team/frontend 0.22.1-dev33 → 0.22.1-dev38

package/src/plugins/impl/plotly/PlotlyPlugin.tsx

@@ -14,11 +14,16 @@ import useEvent from "react-use-event-hook";
 import { useDeepCompareMemoize } from "@/hooks/useDeepCompareMemoize";
 import { useScript } from "@/hooks/useScript";
 import { Arrays } from "@/utils/arrays";
-import { Objects } from "@/utils/objects";
 import {
-  extractClickSelection,
   extractIndices,
   extractPoints,
+  extractSunburstPoints,
+  extractTreemapPoints,
+  hasPureLineTrace,
+  lineSelectionButtons,
+  type ModeBarButton,
+  mergeModeBarButtonsToAdd,
+  shouldHandleClickSelection,
 } from "./selection";
 import { usePlotlyLayout } from "./usePlotlyLayout";
 
@@ -35,6 +40,10 @@ type T =
         x?: number[];
         y?: number[];
       };
+      lasso?: {
+        x?: unknown[];
+        y?: unknown[];
+      };
       // These are kept in the state to persist selections across re-renders
       // on the frontend, but likely not used in the backend.
       selections?: unknown[];
@@ -77,23 +86,6 @@ const LazyPlot = lazy(() =>
   import("./Plot").then((mod) => ({ default: mod.Plot })),
 );
 
-const SUNBURST_DATA_KEYS: (keyof Plotly.SunburstPlotDatum)[] = [
-  "color",
-  "curveNumber",
-  "entry",
-  "hovertext",
-  "id",
-  "label",
-  "parent",
-  "percentEntry",
-  "percentParent",
-  "percentRoot",
-  "pointNumber",
-  "root",
-  "value",
-] as const;
-const TREE_MAP_DATA_KEYS = SUNBURST_DATA_KEYS;
-
 export const PlotlyComponent = memo(
   ({ figure: originalFigure, value, setValue, config }: PlotlyPluginProps) => {
     // Used for rendering LaTeX. TODO: Serve this library from Marimo
@@ -102,7 +94,7 @@ export const PlotlyComponent = memo(
     );
     const isScriptLoaded = scriptStatus === "ready";
 
-    const { figure, layout, handleReset } = usePlotlyLayout({
+    const { figure, layout, setLayout, handleReset } = usePlotlyLayout({
       originalFigure,
       initialValue: value,
       isScriptLoaded,
@@ -112,31 +104,48 @@ export const PlotlyComponent = memo(
       handleReset();
       setValue({});
     });
+    const handleSetDragmode = useEvent(
+      (dragmode: Plotly.Layout["dragmode"]) => {
+        setLayout((prev) => ({ ...prev, dragmode }));
+        setValue((prev) => ({ ...prev, dragmode }));
+      },
+    );
 
     const configMemo = useDeepCompareMemoize(config);
     const plotlyConfig = useMemo((): Partial<Plotly.Config> => {
-      return {
-        displaylogo: false,
-        modeBarButtonsToAdd: [
-          // Custom button to reset the state
-          {
-            name: "reset",
-            title: "Reset state",
-            icon: {
-              svg: `
+      const hasPureLine = hasPureLineTrace(figure.data);
+      const defaultButtons: ModeBarButton[] = [
+        // Custom button to reset the state
+        {
+          name: "reset",
+          title: "Reset state",
+          icon: {
+            svg: `
               <svg xmlns="http://www.w3.org/2000/svg" width="24" height="24" viewBox="0 0 24 24" fill="none" stroke="currentColor"
                 stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-rotate-ccw">
                 <path d="M3 12a9 9 0 1 0 9-9 9.75 9.75 0 0 0-6.74 2.74L3 8" />
                 <path d="M3 3v5h5" />
               </svg>`,
-            },
-            click: handleResetWithClear,
           },
-        ],
+          click: handleResetWithClear,
+        },
+      ];
+      if (hasPureLine) {
+        defaultButtons.push(...lineSelectionButtons(handleSetDragmode));
+      }
+
+      return {
+        displaylogo: false,
         // Prioritize user's config
         ...configMemo,
+        modeBarButtonsToAdd: mergeModeBarButtonsToAdd(
+          defaultButtons,
+          configMemo.modeBarButtonsToAdd as
+            | readonly ModeBarButton[]
+            | undefined,
+        ),
       };
-    }, [handleResetWithClear, configMemo]);
+    }, [handleResetWithClear, handleSetDragmode, configMemo, figure.data]);
 
     return (
       <LazyPlot
@@ -171,6 +180,7 @@ export const PlotlyComponent = memo(
               points: Arrays.EMPTY,
               indices: Arrays.EMPTY,
               range: undefined,
+              lasso: undefined,
             };
           });
         })}
@@ -181,9 +191,7 @@ export const PlotlyComponent = memo(
 
           setValue((prev) => ({
             ...prev,
-            points: evt.points.map((point) =>
-              Objects.pick(point, TREE_MAP_DATA_KEYS),
-            ),
+            points: extractTreemapPoints(evt.points),
           }));
         })}
         onSunburstClick={useEvent((evt: Readonly<Plotly.PlotMouseEvent>) => {
@@ -193,9 +201,7 @@ export const PlotlyComponent = memo(
 
           setValue((prev) => ({
             ...prev,
-            points: evt.points.map((point) =>
-              Objects.pick(point, SUNBURST_DATA_KEYS),
-            ),
+            points: extractSunburstPoints(evt.points),
           }));
         })}
         config={plotlyConfig}
@@ -203,13 +209,21 @@ export const PlotlyComponent = memo(
           if (!evt) {
             return;
           }
-
-          const clickSelection = extractClickSelection(evt);
-          if (!clickSelection) {
+          // Handle clicks for chart types where box/lasso selection
+          // is limited or unavailable (e.g. bar, heatmaps, histograms, pure line traces).
+          if (!shouldHandleClickSelection(evt.points)) {
             return;
           }
-
-          setValue((prev) => ({ ...prev, ...clickSelection }));
+          const extractedPoints = extractPoints(evt.points);
+          const extractedIndices = extractIndices(evt.points);
+          setValue((prev) => ({
+            ...prev,
+            selections: Arrays.EMPTY,
+            range: undefined,
+            lasso: undefined,
+            points: extractedPoints,
+            indices: extractedIndices,
+          }));
         })}
         onSelected={useEvent((evt: Readonly<Plotly.PlotSelectionEvent>) => {
           if (!evt) {
@@ -223,6 +237,10 @@ export const PlotlyComponent = memo(
             points: extractPoints(evt.points),
             indices: extractIndices(evt.points),
             range: evt.range,
+            lasso:
+              "lassoPoints" in evt
+                ? (evt.lassoPoints as { x?: unknown[]; y?: unknown[] })
+                : undefined,
           }));
         })}
         className="w-full"

package/src/plugins/impl/plotly/__tests__/PlotlyPlugin.test.tsx

@@ -0,0 +1,114 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, render, waitFor } from "@testing-library/react";
+import { Suspense } from "react";
+import { describe, expect, it, vi } from "vitest";
+import { SetupMocks } from "@/__mocks__/common";
+import type { Setter } from "@/plugins/types";
+import { PlotlyComponent } from "../PlotlyPlugin";
+
+SetupMocks.resizeObserver();
+
+type CapturedPlotProps = {
+  onClick?: (event: {
+    points: {
+      data?: { type?: string };
+      x?: string | number;
+      y?: string | number;
+      pointIndex?: number;
+      pointNumber?: number;
+      curveNumber?: number;
+    }[];
+  }) => void;
+} | null;
+
+let capturedPlotProps: CapturedPlotProps = null;
+
+vi.mock("../Plot", () => ({
+  Plot: (props: CapturedPlotProps) => {
+    capturedPlotProps = props;
+    return <div data-testid="plotly-mock" />;
+  },
+}));
+
+vi.mock("../usePlotlyLayout", () => ({
+  usePlotlyLayout: ({
+    originalFigure,
+  }: {
+    originalFigure: {
+      data: unknown[];
+      layout: Record<string, unknown>;
+      frames: unknown[] | null;
+    };
+  }) => ({
+    figure: originalFigure,
+    layout: originalFigure.layout,
+    handleReset: vi.fn(),
+  }),
+}));
+
+vi.mock("@/hooks/useScript", () => ({
+  useScript: () => "ready",
+}));
+
+vi.mock("react-use-event-hook", () => ({
+  default: <T,>(callback: T) => callback,
+}));
+
+describe("PlotlyPlugin", () => {
+  it("clicking a bar selects that bar", async () => {
+    const setValue = vi.fn<Setter<unknown>>();
+
+    render(
+      <Suspense fallback={null}>
+        <PlotlyComponent
+          figure={{
+            data: [{ type: "bar" }],
+            layout: {},
+            frames: null,
+          }}
+          value={undefined}
+          setValue={setValue}
+          host={document.createElement("div")}
+          config={{}}
+        />
+      </Suspense>,
+    );
+
+    await waitFor(() => {
+      expect(capturedPlotProps).not.toBeNull();
+    });
+
+    act(() => {
+      capturedPlotProps?.onClick?.({
+        points: [
+          {
+            data: { type: "bar" },
+            x: "Feb",
+            y: 18,
+            pointIndex: 1,
+            pointNumber: 1,
+            curveNumber: 0,
+          },
+        ],
+      });
+    });
+
+    expect(setValue).toHaveBeenCalledTimes(1);
+    const updater = setValue.mock.calls[0][0] as (value: unknown) => unknown;
+    expect(updater({})).toEqual({
+      selections: [],
+      points: [
+        {
+          x: "Feb",
+          y: 18,
+          curveNumber: 0,
+          pointNumber: 1,
+          pointIndex: 1,
+        },
+      ],
+      indices: [1],
+      range: undefined,
+    });
+  });
+});