@marimo-team/islands 0.19.7-dev2 → 0.19.7-dev20

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

@@ -0,0 +1,113 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { Figure } from "react-plotly.js";
+import { describe, expect, it } from "vitest";
+import {
+  computeLayoutOnFigureChange,
+  computeLayoutUpdate,
+  computeOmitKeys,
+  createInitialLayout,
+} from "../usePlotlyLayout";
+
+function createFigure(layoutOverrides: Partial<Plotly.Layout> = {}): Figure {
+  return {
+    data: [],
+    layout: { ...layoutOverrides } as Plotly.Layout,
+    frames: null,
+  };
+}
+
+describe("createInitialLayout", () => {
+  it("sets defaults and merges figure layout", () => {
+    const figure = createFigure({ title: { text: "Test" }, dragmode: "zoom" });
+    const result = createInitialLayout(figure);
+
+    expect(result.autosize).toBe(true);
+    expect(result.height).toBe(540);
+    expect(result.dragmode).toBe("zoom"); // figure overrides default
+    expect(result.title).toEqual({ text: "Test" });
+  });
+
+  it("disables autosize when width is specified", () => {
+    const result = createInitialLayout(createFigure({ width: 800 }));
+    expect(result.autosize).toBe(false);
+  });
+});
+
+describe("computeLayoutOnFigureChange", () => {
+  it("preserves only dragmode/xaxis/yaxis from previous layout (#7964)", () => {
+    const nextFigure = createFigure({ title: { text: "New" } });
+    const prevLayout: Partial<Plotly.Layout> = {
+      dragmode: "zoom",
+      xaxis: { range: [0, 10] },
+      yaxis: { range: [0, 100] },
+      shapes: [{ type: "rect", x0: 0, x1: 1, y0: 0, y1: 1 }],
+      annotations: [{ text: "Old", x: 0, y: 0 }],
+    };
+
+    const result = computeLayoutOnFigureChange(nextFigure, prevLayout);
+
+    // Preserved from prev
+    expect(result.dragmode).toBe("zoom");
+    expect(result.xaxis).toEqual({ range: [0, 10] });
+    expect(result.yaxis).toEqual({ range: [0, 100] });
+    // From new figure
+    expect(result.title).toEqual({ text: "New" });
+    // NOT preserved (the bug fix)
+    expect(result.shapes).toBeUndefined();
+    expect(result.annotations).toBeUndefined();
+  });
+
+  it("uses shapes from new figure, not previous layout", () => {
+    const nextFigure = createFigure({
+      shapes: [{ type: "circle", x0: 0, x1: 1, y0: 0, y1: 1 }],
+    });
+    const prevLayout: Partial<Plotly.Layout> = {
+      shapes: [{ type: "rect", x0: 0, x1: 1, y0: 0, y1: 1 }],
+    };
+
+    const result = computeLayoutOnFigureChange(nextFigure, prevLayout);
+
+    expect(result.shapes).toHaveLength(1);
+    expect(result.shapes?.[0].type).toBe("circle");
+  });
+});
+
+describe("computeOmitKeys", () => {
+  it("omits user-interaction keys unless they changed in figure", () => {
+    const unchanged = computeOmitKeys({}, {});
+    expect([...unchanged]).toEqual(
+      expect.arrayContaining(["autosize", "dragmode", "xaxis", "yaxis"]),
+    );
+
+    const changed = computeOmitKeys(
+      { dragmode: "zoom", xaxis: { range: [0, 10] } },
+      { dragmode: "select", xaxis: { range: [0, 5] } },
+    );
+    expect(changed.has("dragmode")).toBe(false);
+    expect(changed.has("xaxis")).toBe(false);
+    expect(changed.has("autosize")).toBe(true);
+  });
+});
+
+describe("computeLayoutUpdate", () => {
+  it("merges figure layout while respecting omit keys", () => {
+    // dragmode unchanged in figure -> preserve prev layout's dragmode
+    const result1 = computeLayoutUpdate(
+      { dragmode: "pan", title: { text: "New" } },
+      { dragmode: "pan" },
+      { dragmode: "zoom", height: 400 },
+    );
+    expect(result1.dragmode).toBe("zoom");
+    expect(result1.title).toEqual({ text: "New" });
+    expect(result1.height).toBe(400);
+
+    // dragmode changed in figure -> use figure's dragmode
+    const result2 = computeLayoutUpdate(
+      { dragmode: "pan" },
+      { dragmode: "select" },
+      { dragmode: "zoom" },
+    );
+    expect(result2.dragmode).toBe("pan");
+  });
+});

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

@@ -0,0 +1,163 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { usePrevious } from "@uidotdev/usehooks";
+import { isEqual, pick } from "lodash-es";
+import { useEffect, useState } from "react";
+import type { Figure } from "react-plotly.js";
+import { Objects } from "@/utils/objects";
+
+/**
+ * Keys that are preserved across figure updates when set by user interaction.
+ * These include dragmode and axis settings that users may adjust.
+ */
+export const PERSISTED_LAYOUT_KEYS = ["dragmode", "xaxis", "yaxis"] as const;
+
+/**
+ * Keys that are omitted from layout updates unless they changed in the figure.
+ * This prevents overwriting user interactions like zoom/pan.
+ */
+export const LAYOUT_OMIT_KEYS: (keyof Plotly.Layout)[] = [
+  "autosize",
+  "dragmode",
+  "xaxis",
+  "yaxis",
+];
+
+/**
+ * Creates the initial layout for a Plotly figure with sensible defaults.
+ */
+export function createInitialLayout(figure: Figure): Partial<Plotly.Layout> {
+  // Enable autosize if width is not specified
+  const shouldAutoSize = figure.layout.width === undefined;
+  return {
+    autosize: shouldAutoSize,
+    dragmode: "select",
+    height: 540,
+    // Prioritize user's config
+    ...figure.layout,
+  };
+}
+
+/**
+ * Computes the updated layout when the figure changes.
+ * Preserves user-interaction values (dragmode, xaxis, yaxis) while
+ * taking everything else from the new figure's layout.
+ */
+export function computeLayoutOnFigureChange(
+  nextFigure: Figure,
+  prevLayout: Partial<Plotly.Layout>,
+): Partial<Plotly.Layout> {
+  return {
+    ...createInitialLayout(nextFigure),
+    ...pick(prevLayout, PERSISTED_LAYOUT_KEYS),
+  };
+}
+
+/**
+ * Computes which keys to omit from layout updates based on what changed.
+ * If a key changed in the figure, we should update it even if it's normally omitted.
+ */
+export function computeOmitKeys(
+  currentLayout: Partial<Plotly.Layout>,
+  previousLayout: Partial<Plotly.Layout>,
+): Set<keyof Plotly.Layout> {
+  const omitKeys = new Set<keyof Plotly.Layout>(LAYOUT_OMIT_KEYS);
+
+  // If the key was updated externally (e.g. can be specifically passed in the config)
+  // then we need to update the layout
+  for (const key of omitKeys) {
+    if (!isEqual(currentLayout[key], previousLayout[key])) {
+      omitKeys.delete(key);
+    }
+  }
+
+  return omitKeys;
+}
+
+/**
+ * Computes the layout update when figure.layout changes.
+ * Omits keys that shouldn't override user interactions unless they changed.
+ */
+export function computeLayoutUpdate(
+  figureLayout: Partial<Plotly.Layout>,
+  previousFigureLayout: Partial<Plotly.Layout>,
+  prevLayout: Partial<Plotly.Layout>,
+): Partial<Plotly.Layout> {
+  const omitKeys = computeOmitKeys(figureLayout, previousFigureLayout);
+  const layoutUpdate = Objects.omit(figureLayout, omitKeys);
+  return { ...prevLayout, ...layoutUpdate };
+}
+
+interface UsePlotlyLayoutOptions {
+  originalFigure: Figure;
+  initialValue?: Partial<Plotly.Layout>;
+  isScriptLoaded?: boolean;
+}
+
+interface UsePlotlyLayoutResult {
+  figure: Figure;
+  layout: Partial<Plotly.Layout>;
+  setLayout: React.Dispatch<React.SetStateAction<Partial<Plotly.Layout>>>;
+  handleReset: () => void;
+}
+
+/**
+ * Hook that manages the Plotly figure and layout state.
+ *
+ * This hook handles:
+ * - Cloning the figure to prevent Plotly mutations
+ * - Managing layout state with proper preservation of user interactions
+ * - Syncing layout when the figure changes
+ * - Providing a reset function to restore original state
+ */
+export function usePlotlyLayout({
+  originalFigure,
+  initialValue,
+  isScriptLoaded = true,
+}: UsePlotlyLayoutOptions): UsePlotlyLayoutResult {
+  const [figure, setFigure] = useState(() => {
+    // We clone the figure since Plotly mutates the figure in place
+    return structuredClone(originalFigure);
+  });
+
+  const [layout, setLayout] = useState<Partial<Plotly.Layout>>(() => {
+    return {
+      ...createInitialLayout(figure),
+      // Override with persisted values (dragmode, xaxis, yaxis)
+      ...initialValue,
+    };
+  });
+
+  // Update figure and layout when originalFigure changes
+  useEffect(() => {
+    const nextFigure = structuredClone(originalFigure);
+    setFigure(nextFigure);
+    // Start with the new figure's layout, then only preserve user-interaction
+    // values (dragmode, xaxis, yaxis) from the previous layout.
+    // We don't want to preserve other properties like `shapes` from the previous
+    // layout, as they should be fully controlled by the figure prop.
+    setLayout((prev) => computeLayoutOnFigureChange(nextFigure, prev));
+  }, [originalFigure, isScriptLoaded]);
+
+  const prevFigure = usePrevious(figure) ?? figure;
+
+  // Sync layout when figure.layout changes
+  useEffect(() => {
+    setLayout((prev) =>
+      computeLayoutUpdate(figure.layout, prevFigure.layout, prev),
+    );
+  }, [figure.layout, prevFigure.layout]);
+
+  const handleReset = () => {
+    const nextFigure = structuredClone(originalFigure);
+    setFigure(nextFigure);
+    setLayout(createInitialLayout(nextFigure));
+  };
+
+  return {
+    figure,
+    layout,
+    setLayout,
+    handleReset,
+  };
+}

package/src/utils/__tests__/download.test.tsx

@@ -155,10 +155,16 @@ describe("getImageDataUrlForCell", () => {
     const result = await getImageDataUrlForCell("cell-1" as CellId);
 
     expect(result).toBe(mockDataUrl);
-    expect(toPng).toHaveBeenCalledWith(mockElement);
+    expect(toPng).toHaveBeenCalledWith(
+      mockElement,
+      expect.objectContaining({
+        filter: expect.any(Function),
+        onImageErrorHandler: expect.any(Function),
+      }),
+    );
   });
 
-  it("should add printing classes before capture", async () => {
+  it("should add printing classes before capture when enablePrintMode is true", async () => {
     vi.mocked(toPng).mockImplementation(async () => {
       // Check classes are applied during capture
       expect(mockElement.classList.contains("printing-output")).toBe(true);
@@ -167,16 +173,40 @@ describe("getImageDataUrlForCell", () => {
       return mockDataUrl;
     });
 
-    await getImageDataUrlForCell("cell-1" as CellId);
+    await getImageDataUrlForCell("cell-1" as CellId, true);
   });
 
-  it("should remove printing classes after capture", async () => {
+  it("should remove printing classes after capture when enablePrintMode is true", async () => {
     vi.mocked(toPng).mockResolvedValue(mockDataUrl);
 
-    await getImageDataUrlForCell("cell-1" as CellId);
+    await getImageDataUrlForCell("cell-1" as CellId, true);
+
+    expect(mockElement.classList.contains("printing-output")).toBe(false);
+    expect(document.body.classList.contains("printing")).toBe(false);
+  });
+
+  it("should add printing-output but NOT body.printing when enablePrintMode is false", async () => {
+    vi.mocked(toPng).mockImplementation(async () => {
+      // printing-output should still be added to the element
+      expect(mockElement.classList.contains("printing-output")).toBe(true);
+      // but body.printing should NOT be added
+      expect(document.body.classList.contains("printing")).toBe(false);
+      expect(mockElement.style.overflow).toBe("auto");
+      return mockDataUrl;
+    });
+
+    await getImageDataUrlForCell("cell-1" as CellId, false);
+  });
+
+  it("should cleanup printing-output when enablePrintMode is false", async () => {
+    mockElement.style.overflow = "hidden";
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    await getImageDataUrlForCell("cell-1" as CellId, false);
 
     expect(mockElement.classList.contains("printing-output")).toBe(false);
     expect(document.body.classList.contains("printing")).toBe(false);
+    expect(mockElement.style.overflow).toBe("hidden");
   });
 
   it("should restore original overflow style after capture", async () => {
@@ -207,7 +237,7 @@ describe("getImageDataUrlForCell", () => {
     expect(mockElement.style.overflow).toBe("scroll");
   });
 
-  it("should maintain body.printing during concurrent captures", async () => {
+  it("should maintain body.printing during concurrent captures when enablePrintMode is true", async () => {
     // Create a second element
     const mockElement2 = document.createElement("div");
     mockElement2.id = CellOutputId.create("cell-2" as CellId);
@@ -241,9 +271,9 @@ describe("getImageDataUrlForCell", () => {
       return mockDataUrl;
     });
 
-    // Start both captures concurrently
-    const capture1 = getImageDataUrlForCell("cell-1" as CellId);
-    const capture2 = getImageDataUrlForCell("cell-2" as CellId);
+    // Start both captures concurrently with enablePrintMode = true
+    const capture1 = getImageDataUrlForCell("cell-1" as CellId, true);
+    const capture2 = getImageDataUrlForCell("cell-2" as CellId, true);
 
     // Let second capture complete first
     resolveSecond!();
@@ -264,6 +294,30 @@ describe("getImageDataUrlForCell", () => {
 
     mockElement2.remove();
   });
+
+  it("should not interfere with body.printing during concurrent captures when enablePrintMode is false", async () => {
+    // Create a second element
+    const mockElement2 = document.createElement("div");
+    mockElement2.id = CellOutputId.create("cell-2" as CellId);
+    document.body.append(mockElement2);
+
+    vi.mocked(toPng).mockImplementation(async () => {
+      // body.printing should never be added when enablePrintMode is false
+      expect(document.body.classList.contains("printing")).toBe(false);
+      return mockDataUrl;
+    });
+
+    // Start both captures concurrently with enablePrintMode = false
+    const capture1 = getImageDataUrlForCell("cell-1" as CellId, false);
+    const capture2 = getImageDataUrlForCell("cell-2" as CellId, false);
+
+    await Promise.all([capture1, capture2]);
+
+    // body.printing should still not be present
+    expect(document.body.classList.contains("printing")).toBe(false);
+
+    mockElement2.remove();
+  });
 });
 
 describe("downloadHTMLAsImage", () => {
@@ -304,7 +358,13 @@ describe("downloadHTMLAsImage", () => {
 
     await downloadHTMLAsImage({ element: mockElement, filename: "test" });
 
-    expect(toPng).toHaveBeenCalledWith(mockElement);
+    expect(toPng).toHaveBeenCalledWith(
+      mockElement,
+      expect.objectContaining({
+        filter: expect.any(Function),
+        onImageErrorHandler: expect.any(Function),
+      }),
+    );
     expect(mockAnchor.href).toBe(mockDataUrl);
     expect(mockAnchor.download).toBe("test.png");
     expect(mockAnchor.click).toHaveBeenCalled();
@@ -342,7 +402,7 @@ describe("downloadHTMLAsImage", () => {
     expect(cleanup).toHaveBeenCalled();
   });
 
-  it("should not add body.printing when prepare is provided", async () => {
+  it("should delegate body.printing management to prepare function", async () => {
     let bodyPrintingDuringCapture = false;
     vi.mocked(toPng).mockImplementation(async () => {
       // Capture the state during toPng execution
@@ -350,7 +410,14 @@ describe("downloadHTMLAsImage", () => {
       return mockDataUrl;
     });
     const cleanup = vi.fn();
-    const prepare = vi.fn().mockReturnValue(cleanup);
+    // Mock prepare that adds body.printing
+    const prepare = vi.fn().mockImplementation(() => {
+      document.body.classList.add("printing");
+      return () => {
+        document.body.classList.remove("printing");
+        cleanup();
+      };
+    });
 
     await downloadHTMLAsImage({
       element: mockElement,
@@ -358,9 +425,8 @@ describe("downloadHTMLAsImage", () => {
       prepare,
     });
 
-    // body.printing should NOT be added by downloadHTMLAsImage when prepare is provided
-    // (the prepare function is responsible for managing its own classes)
-    expect(bodyPrintingDuringCapture).toBe(false);
+    // body.printing should be added by prepare function
+    expect(bodyPrintingDuringCapture).toBe(true);
     expect(document.body.classList.contains("printing")).toBe(false);
     expect(prepare).toHaveBeenCalledWith(mockElement);
     expect(cleanup).toHaveBeenCalled();
@@ -434,7 +500,13 @@ describe("downloadCellOutputAsImage", () => {
 
     await downloadCellOutputAsImage("cell-1" as CellId, "result");
 
-    expect(toPng).toHaveBeenCalledWith(mockElement);
+    expect(toPng).toHaveBeenCalledWith(
+      mockElement,
+      expect.objectContaining({
+        filter: expect.any(Function),
+        onImageErrorHandler: expect.any(Function),
+      }),
+    );
     expect(mockAnchor.download).toBe("result.png");
   });
 

package/src/utils/download.ts

@@ -1,11 +1,12 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { toPng } from "html-to-image";
+
 import { toast } from "@/components/ui/use-toast";
 import { type CellId, CellOutputId } from "@/core/cells/ids";
 import { getRequestClient } from "@/core/network/requests";
 import { Filenames } from "@/utils/filenames";
 import { Paths } from "@/utils/paths";
 import { prettyError } from "./errors";
+import { toPng } from "./html-to-image";
 import { Logger } from "./Logger";
 
 /**
@@ -59,34 +60,51 @@ function releaseBodyPrinting() {
   }
 }
 
-/*
+/**
  * Prepare a cell element for screenshot capture.
- * Returns a cleanup function that should be called when the screenshot is complete.
+ *
+ * @param element - The cell output element to prepare
+ * @param enablePrintMode - When true, adds a 'printing' class to the body.
+ *   This can cause layout shifts that cause the page to scroll.
+ * @returns A cleanup function to restore the element's original state
  */
-function prepareCellElementForScreenshot(element: HTMLElement) {
+function prepareCellElementForScreenshot(
+  element: HTMLElement,
+  enablePrintMode: boolean,
+) {
   element.classList.add("printing-output");
-  acquireBodyPrinting();
+  if (enablePrintMode) {
+    acquireBodyPrinting();
+  }
   const originalOverflow = element.style.overflow;
   element.style.overflow = "auto";
 
   return () => {
     element.classList.remove("printing-output");
-    releaseBodyPrinting();
+    if (enablePrintMode) {
+      releaseBodyPrinting();
+    }
     element.style.overflow = originalOverflow;
   };
 }
 
 /**
  * Capture a cell output as a PNG data URL.
+ *
+ * @param cellId - The ID of the cell to capture
+ * @param enablePrintMode - When true, enables print mode which adds a 'printing' class to the body.
+ *   This can cause layout shifts that cause the page to scroll.
+ * @returns The PNG as a data URL, or undefined if the cell element wasn't found
  */
 export async function getImageDataUrlForCell(
   cellId: CellId,
+  enablePrintMode = true,
 ): Promise<string | undefined> {
   const element = findElementForCell(cellId);
   if (!element) {
     return;
   }
-  const cleanup = prepareCellElementForScreenshot(element);
+  const cleanup = prepareCellElementForScreenshot(element, enablePrintMode);
 
   try {
     return await toPng(element);
@@ -110,7 +128,7 @@ export async function downloadCellOutputAsImage(
   await downloadHTMLAsImage({
     element,
     filename,
-    prepare: prepareCellElementForScreenshot,
+    prepare: () => prepareCellElementForScreenshot(element, true),
   });
 }
 
@@ -127,7 +145,6 @@ export async function downloadHTMLAsImage(opts: {
 
   let cleanup: (() => void) | undefined;
   if (prepare) {
-    // Let the prepare function handle adding classes (e.g., body.printing)
     cleanup = prepare(element);
   } else {
     // When no prepare function is provided (e.g., downloading full notebook),

package/src/utils/html-to-image.ts

@@ -0,0 +1,41 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { toPng as htmlToImageToPng } from "html-to-image";
+import { Logger } from "./Logger";
+
+export type HtmlToImageOptions = Parameters<typeof htmlToImageToPng>[1];
+
+/**
+ * Default options for html-to-image conversions.
+ * These handle common edge cases like filtering out toolbars and logging errors.
+ */
+export const defaultHtmlToImageOptions: HtmlToImageOptions = {
+  filter: (node) => {
+    try {
+      if ("classList" in node) {
+        // Filter out matplotlib toolbars
+        return !node.classList.contains("mpl-toolbar");
+      }
+      return true;
+    } catch (error) {
+      Logger.error("Error filtering node:", error);
+      return true;
+    }
+  },
+  onImageErrorHandler: (event) => {
+    Logger.error("Error loading image:", event);
+  },
+};
+
+/**
+ * Convert an HTML element to a PNG data URL.
+ * This is a wrapper around html-to-image's toPng with default options applied.
+ */
+export function toPng(
+  element: HTMLElement,
+  options?: HtmlToImageOptions,
+): Promise<string> {
+  return htmlToImageToPng(element, {
+    ...defaultHtmlToImageOptions,
+    ...options,
+  });
+}