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

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/__tests__/iframe.test.ts

@@ -0,0 +1,139 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { captureIframeAsImage } from "../iframe";
+
+// Mock html-to-image
+vi.mock("html-to-image", () => ({
+  toPng: vi.fn().mockResolvedValue("data:image/png;base64,mock"),
+}));
+
+describe("captureIframeAsImage", () => {
+  const originalCreateElement = document.createElement.bind(document);
+
+  beforeEach(() => {
+    // Mock devicePixelRatio
+    Object.defineProperty(window, "devicePixelRatio", {
+      value: 1,
+      writable: true,
+    });
+
+    // Mock canvas for placeholder generation
+    vi.spyOn(document, "createElement").mockImplementation((tagName) => {
+      const element = originalCreateElement(tagName);
+      if (tagName === "canvas") {
+        const mockCtx = {
+          scale: vi.fn(),
+          fillStyle: "",
+          fillRect: vi.fn(),
+          strokeStyle: "",
+          strokeRect: vi.fn(),
+          font: "",
+          textAlign: "",
+          textBaseline: "",
+          fillText: vi.fn(),
+          measureText: vi.fn().mockReturnValue({ width: 10 }),
+        };
+        vi.spyOn(element as HTMLCanvasElement, "getContext").mockReturnValue(
+          mockCtx as unknown as CanvasRenderingContext2D,
+        );
+        vi.spyOn(element as HTMLCanvasElement, "toDataURL").mockReturnValue(
+          "data:image/png;base64,placeholder",
+        );
+      }
+      return element;
+    });
+  });
+
+  afterEach(() => {
+    vi.clearAllMocks();
+    vi.restoreAllMocks();
+  });
+
+  it("should return null when element has no iframe", async () => {
+    const element = document.createElement("div");
+    element.innerHTML = "<p>No iframe here</p>";
+
+    const result = await captureIframeAsImage(element);
+    expect(result).toBeNull();
+  });
+
+  it("should return placeholder for external iframe", async () => {
+    const element = document.createElement("div");
+    element.innerHTML = '<iframe src="https://external.com/page"></iframe>';
+
+    const result = await captureIframeAsImage(element);
+
+    expect(result).not.toBeNull();
+    expect(result).toMatch(/^data:image\/png;base64,/);
+  });
+
+  it("should return placeholder for cross-origin iframe", async () => {
+    const element = document.createElement("div");
+    element.innerHTML =
+      '<iframe src="https://www.openstreetmap.org/export/embed.html"></iframe>';
+
+    const result = await captureIframeAsImage(element);
+
+    expect(result).not.toBeNull();
+    expect(result).toMatch(/^data:image\/png;base64,/);
+  });
+
+  it("should return null for about:blank iframe without body", async () => {
+    const element = document.createElement("div");
+    const iframe = document.createElement("iframe");
+    iframe.src = "about:blank";
+    element.append(iframe);
+
+    // The iframe has no body accessible in jsdom
+    const result = await captureIframeAsImage(element);
+
+    // In jsdom, contentDocument may not be accessible
+    expect(result).toBeNull();
+  });
+
+  it("should handle iframe with relative same-origin src", async () => {
+    const element = document.createElement("div");
+    element.innerHTML = '<iframe src="./@file/123.html"></iframe>';
+
+    // Same-origin relative URL, but contentDocument not accessible in jsdom
+    const result = await captureIframeAsImage(element);
+
+    // Returns null because jsdom can't access contentDocument for file URLs
+    expect(result).toBeNull();
+  });
+
+  it("should detect external URL from various formats", async () => {
+    const externalUrls = [
+      "https://example.com",
+      "http://external.org/path",
+      "https://sub.domain.com:8080/page",
+    ];
+
+    for (const url of externalUrls) {
+      const element = document.createElement("div");
+      element.innerHTML = `<iframe src="${url}"></iframe>`;
+
+      const result = await captureIframeAsImage(element);
+
+      expect(result).not.toBeNull();
+      expect(result).toMatch(/^data:image\/png;base64,/);
+    }
+  });
+
+  it("should not treat same-origin URLs as external", async () => {
+    // Same-origin URLs - these should try to capture, not return placeholder
+    const sameOriginUrls = ["/local/path", "./relative/path", "../parent/path"];
+
+    for (const url of sameOriginUrls) {
+      const element = document.createElement("div");
+      element.innerHTML = `<iframe src="${url}"></iframe>`;
+
+      const result = await captureIframeAsImage(element);
+
+      // In jsdom, these return null because contentDocument isn't accessible
+      // but they should NOT return a placeholder (which would indicate external detection)
+      // The key is they don't trigger the external URL path
+      expect(result).toBeNull();
+    }
+  });
+});