@marimo-team/frontend 0.19.5-dev40 → 0.19.5-dev44

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

@@ -1,6 +1,19 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { beforeEach, describe, expect, it, vi } from "vitest";
-import { withLoadingToast } from "../download";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import type { CellId } from "@/core/cells/ids";
+import { CellOutputId } from "@/core/cells/ids";
+import {
+  downloadByURL,
+  downloadCellOutputAsImage,
+  downloadHTMLAsImage,
+  getImageDataUrlForCell,
+  withLoadingToast,
+} from "../download";
+
+// Mock html-to-image
+vi.mock("html-to-image", () => ({
+  toPng: vi.fn(),
+}));
 
 // Mock the toast module
 const mockDismiss = vi.fn();
@@ -15,7 +28,23 @@ vi.mock("@/components/icons/spinner", () => ({
   Spinner: () => "MockSpinner",
 }));
 
+// Mock Logger
+vi.mock("@/utils/Logger", () => ({
+  Logger: {
+    error: vi.fn(),
+  },
+}));
+
+// Mock Filenames
+vi.mock("@/utils/filenames", () => ({
+  Filenames: {
+    toPNG: (name: string) => `${name}.png`,
+  },
+}));
+
+import { toPng } from "html-to-image";
 import { toast } from "@/components/ui/use-toast";
+import { Logger } from "@/utils/Logger";
 
 describe("withLoadingToast", () => {
   beforeEach(() => {
@@ -30,6 +59,7 @@ describe("withLoadingToast", () => {
     expect(toast).toHaveBeenCalledTimes(1);
     expect(toast).toHaveBeenCalledWith({
       title: "Loading...",
+      duration: Infinity,
     });
     expect(mockDismiss).toHaveBeenCalledTimes(1);
     expect(result).toBe("success");
@@ -94,3 +124,369 @@ describe("withLoadingToast", () => {
     expect(mockDismiss).toHaveBeenCalledTimes(1);
   });
 });
+
+describe("getImageDataUrlForCell", () => {
+  const mockDataUrl = "data:image/png;base64,mockbase64data";
+  let mockElement: HTMLElement;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockElement = document.createElement("div");
+    mockElement.id = CellOutputId.create("cell-1" as CellId);
+    document.body.append(mockElement);
+  });
+
+  afterEach(() => {
+    mockElement.remove();
+  });
+
+  it("should return undefined if element is not found", async () => {
+    const result = await getImageDataUrlForCell("nonexistent" as CellId);
+
+    expect(result).toBeUndefined();
+    expect(Logger.error).toHaveBeenCalledWith(
+      "Output element not found for cell nonexistent",
+    );
+  });
+
+  it("should capture screenshot and return data URL", async () => {
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    const result = await getImageDataUrlForCell("cell-1" as CellId);
+
+    expect(result).toBe(mockDataUrl);
+    expect(toPng).toHaveBeenCalledWith(mockElement);
+  });
+
+  it("should add printing classes before capture", async () => {
+    vi.mocked(toPng).mockImplementation(async () => {
+      // Check classes are applied during capture
+      expect(mockElement.classList.contains("printing-output")).toBe(true);
+      expect(document.body.classList.contains("printing")).toBe(true);
+      expect(mockElement.style.overflow).toBe("auto");
+      return mockDataUrl;
+    });
+
+    await getImageDataUrlForCell("cell-1" as CellId);
+  });
+
+  it("should remove printing classes after capture", async () => {
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    await getImageDataUrlForCell("cell-1" as CellId);
+
+    expect(mockElement.classList.contains("printing-output")).toBe(false);
+    expect(document.body.classList.contains("printing")).toBe(false);
+  });
+
+  it("should restore original overflow style after capture", async () => {
+    mockElement.style.overflow = "hidden";
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    await getImageDataUrlForCell("cell-1" as CellId);
+
+    expect(mockElement.style.overflow).toBe("hidden");
+  });
+
+  it("should throw error on failure", async () => {
+    vi.mocked(toPng).mockRejectedValue(new Error("Capture failed"));
+
+    await expect(getImageDataUrlForCell("cell-1" as CellId)).rejects.toThrow(
+      "Capture failed",
+    );
+  });
+
+  it("should cleanup even on failure", async () => {
+    mockElement.style.overflow = "scroll";
+    vi.mocked(toPng).mockRejectedValue(new Error("Capture failed"));
+
+    await expect(getImageDataUrlForCell("cell-1" as CellId)).rejects.toThrow();
+
+    expect(mockElement.classList.contains("printing-output")).toBe(false);
+    expect(document.body.classList.contains("printing")).toBe(false);
+    expect(mockElement.style.overflow).toBe("scroll");
+  });
+
+  it("should maintain body.printing during concurrent captures", async () => {
+    // Create a second element
+    const mockElement2 = document.createElement("div");
+    mockElement2.id = CellOutputId.create("cell-2" as CellId);
+    document.body.append(mockElement2);
+
+    // Track body.printing state during each capture
+    const printingStateDuringCaptures: boolean[] = [];
+    let resolveFirst: () => void;
+    let resolveSecond: () => void;
+
+    const firstPromise = new Promise<void>((resolve) => {
+      resolveFirst = resolve;
+    });
+    const secondPromise = new Promise<void>((resolve) => {
+      resolveSecond = resolve;
+    });
+
+    vi.mocked(toPng).mockImplementation(async (element) => {
+      printingStateDuringCaptures.push(
+        document.body.classList.contains("printing"),
+      );
+
+      // Simulate async work - first capture takes longer
+      await (element.id.includes("cell-1") ? firstPromise : secondPromise);
+
+      // Check state again after waiting
+      printingStateDuringCaptures.push(
+        document.body.classList.contains("printing"),
+      );
+
+      return mockDataUrl;
+    });
+
+    // Start both captures concurrently
+    const capture1 = getImageDataUrlForCell("cell-1" as CellId);
+    const capture2 = getImageDataUrlForCell("cell-2" as CellId);
+
+    // Let second capture complete first
+    resolveSecond!();
+    await new Promise((r) => setTimeout(r, 0));
+
+    // body.printing should still be present because cell-1 is still capturing
+    expect(document.body.classList.contains("printing")).toBe(true);
+
+    // Now let first capture complete
+    resolveFirst!();
+    await Promise.all([capture1, capture2]);
+
+    // After all captures complete, body.printing should be removed
+    expect(document.body.classList.contains("printing")).toBe(false);
+
+    // All captures should have seen body.printing = true
+    expect(printingStateDuringCaptures.every(Boolean)).toBe(true);
+
+    mockElement2.remove();
+  });
+});
+
+describe("downloadHTMLAsImage", () => {
+  const mockDataUrl = "data:image/png;base64,mockbase64data";
+  let mockElement: HTMLElement;
+  let mockAppEl: HTMLElement;
+  let mockAnchor: HTMLAnchorElement;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockElement = document.createElement("div");
+    mockAppEl = document.createElement("div");
+    mockAppEl.id = "App";
+    // Mock scrollTo since jsdom doesn't implement it
+    mockAppEl.scrollTo = vi.fn();
+    document.body.append(mockElement);
+    document.body.append(mockAppEl);
+
+    // Mock anchor element for download
+    mockAnchor = document.createElement("a");
+    vi.spyOn(document, "createElement").mockReturnValue(mockAnchor);
+    vi.spyOn(mockAnchor, "click").mockImplementation(() => {
+      // <noop></noop>
+    });
+    vi.spyOn(mockAnchor, "remove").mockImplementation(() => {
+      // noop
+    });
+  });
+
+  afterEach(() => {
+    mockElement.remove();
+    mockAppEl.remove();
+    vi.restoreAllMocks();
+  });
+
+  it("should download image without prepare function", async () => {
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    await downloadHTMLAsImage({ element: mockElement, filename: "test" });
+
+    expect(toPng).toHaveBeenCalledWith(mockElement);
+    expect(mockAnchor.href).toBe(mockDataUrl);
+    expect(mockAnchor.download).toBe("test.png");
+    expect(mockAnchor.click).toHaveBeenCalled();
+  });
+
+  it("should add body.printing class without prepare function", async () => {
+    vi.mocked(toPng).mockImplementation(async () => {
+      expect(document.body.classList.contains("printing")).toBe(true);
+      return mockDataUrl;
+    });
+
+    await downloadHTMLAsImage({ element: mockElement, filename: "test" });
+  });
+
+  it("should remove body.printing class after download without prepare", async () => {
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    await downloadHTMLAsImage({ element: mockElement, filename: "test" });
+
+    expect(document.body.classList.contains("printing")).toBe(false);
+  });
+
+  it("should use prepare function when provided", async () => {
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+    const cleanup = vi.fn();
+    const prepare = vi.fn().mockReturnValue(cleanup);
+
+    await downloadHTMLAsImage({
+      element: mockElement,
+      filename: "test",
+      prepare,
+    });
+
+    expect(prepare).toHaveBeenCalledWith(mockElement);
+    expect(cleanup).toHaveBeenCalled();
+  });
+
+  it("should not add body.printing when prepare is provided", async () => {
+    let bodyPrintingDuringCapture = false;
+    vi.mocked(toPng).mockImplementation(async () => {
+      // Capture the state during toPng execution
+      bodyPrintingDuringCapture = document.body.classList.contains("printing");
+      return mockDataUrl;
+    });
+    const cleanup = vi.fn();
+    const prepare = vi.fn().mockReturnValue(cleanup);
+
+    await downloadHTMLAsImage({
+      element: mockElement,
+      filename: "test",
+      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);
+    expect(document.body.classList.contains("printing")).toBe(false);
+    expect(prepare).toHaveBeenCalledWith(mockElement);
+    expect(cleanup).toHaveBeenCalled();
+  });
+
+  it("should show error toast on failure", async () => {
+    vi.mocked(toPng).mockRejectedValue(new Error("Failed"));
+
+    await downloadHTMLAsImage({ element: mockElement, filename: "test" });
+
+    expect(toast).toHaveBeenCalledWith({
+      title: "Error",
+      description: "Failed to download as PNG.",
+      variant: "danger",
+    });
+  });
+
+  it("should cleanup on failure", async () => {
+    vi.mocked(toPng).mockRejectedValue(new Error("Failed"));
+
+    await downloadHTMLAsImage({ element: mockElement, filename: "test" });
+
+    expect(document.body.classList.contains("printing")).toBe(false);
+  });
+});
+
+describe("downloadCellOutputAsImage", () => {
+  const mockDataUrl = "data:image/png;base64,mockbase64data";
+  let mockElement: HTMLElement;
+  let mockAppEl: HTMLElement;
+  let mockAnchor: HTMLAnchorElement;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockElement = document.createElement("div");
+    mockElement.id = CellOutputId.create("cell-1" as CellId);
+    mockAppEl = document.createElement("div");
+    mockAppEl.id = "App";
+    // Mock scrollTo since jsdom doesn't implement it
+    mockAppEl.scrollTo = vi.fn();
+    document.body.append(mockElement);
+    document.body.append(mockAppEl);
+
+    mockAnchor = document.createElement("a");
+    vi.spyOn(document, "createElement").mockReturnValue(mockAnchor);
+    vi.spyOn(mockAnchor, "click").mockImplementation(() => {
+      // <noop></noop>
+    });
+    vi.spyOn(mockAnchor, "remove").mockImplementation(() => {
+      // <noop></noop>
+    });
+  });
+
+  afterEach(() => {
+    mockElement.remove();
+    mockAppEl.remove();
+    vi.restoreAllMocks();
+  });
+
+  it("should return early if element not found", async () => {
+    await downloadCellOutputAsImage("nonexistent" as CellId, "test");
+
+    expect(toPng).not.toHaveBeenCalled();
+    expect(Logger.error).toHaveBeenCalledWith(
+      "Output element not found for cell nonexistent",
+    );
+  });
+
+  it("should download cell output as image", async () => {
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    await downloadCellOutputAsImage("cell-1" as CellId, "result");
+
+    expect(toPng).toHaveBeenCalledWith(mockElement);
+    expect(mockAnchor.download).toBe("result.png");
+  });
+
+  it("should apply cell-specific preparation", async () => {
+    vi.mocked(toPng).mockImplementation(async () => {
+      // Check that cell-specific classes are applied
+      expect(mockElement.classList.contains("printing-output")).toBe(true);
+      expect(document.body.classList.contains("printing")).toBe(true);
+      expect(mockElement.style.overflow).toBe("auto");
+      return mockDataUrl;
+    });
+
+    await downloadCellOutputAsImage("cell-1" as CellId, "result");
+  });
+
+  it("should cleanup after download", async () => {
+    mockElement.style.overflow = "visible";
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    await downloadCellOutputAsImage("cell-1" as CellId, "result");
+
+    expect(mockElement.classList.contains("printing-output")).toBe(false);
+    expect(document.body.classList.contains("printing")).toBe(false);
+    expect(mockElement.style.overflow).toBe("visible");
+  });
+});
+
+describe("downloadByURL", () => {
+  let mockAnchor: HTMLAnchorElement;
+
+  beforeEach(() => {
+    mockAnchor = document.createElement("a");
+    vi.spyOn(document, "createElement").mockReturnValue(mockAnchor);
+    vi.spyOn(mockAnchor, "click").mockImplementation(() => {
+      // <noop></noop>
+    });
+    vi.spyOn(mockAnchor, "remove").mockImplementation(() => {
+      // <noop></noop>
+    });
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it("should create anchor, set attributes, click, and remove", () => {
+    downloadByURL("data:test", "filename.png");
+
+    expect(document.createElement).toHaveBeenCalledWith("a");
+    expect(mockAnchor.href).toBe("data:test");
+    expect(mockAnchor.download).toBe("filename.png");
+    expect(mockAnchor.click).toHaveBeenCalled();
+    expect(mockAnchor.remove).toHaveBeenCalled();
+  });
+});

package/src/utils/download.ts

@@ -1,10 +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 { Logger } from "./Logger";
 
 /**
  * Show a loading toast while an async operation is in progress.
@@ -16,6 +18,7 @@ export async function withLoadingToast<T>(
 ): Promise<T> {
   const loadingToast = toast({
     title,
+    duration: Infinity,
   });
   try {
     const result = await fn();
@@ -27,15 +30,111 @@ export async function withLoadingToast<T>(
   }
 }
 
-export async function downloadHTMLAsImage(
-  element: HTMLElement,
+function findElementForCell(cellId: CellId): HTMLElement | undefined {
+  const element = document.getElementById(CellOutputId.create(cellId));
+  if (!element) {
+    Logger.error(`Output element not found for cell ${cellId}`);
+    return;
+  }
+  return element;
+}
+
+/**
+ * Reference counter for body.printing class to handle concurrent screenshot captures.
+ * Only adds the class when count goes 0→1, only removes when count goes 1→0.
+ */
+let bodyPrintingRefCount = 0;
+
+function acquireBodyPrinting() {
+  bodyPrintingRefCount++;
+  if (bodyPrintingRefCount === 1) {
+    document.body.classList.add("printing");
+  }
+}
+
+function releaseBodyPrinting() {
+  bodyPrintingRefCount--;
+  if (bodyPrintingRefCount === 0) {
+    document.body.classList.remove("printing");
+  }
+}
+
+/*
+ * Prepare a cell element for screenshot capture.
+ * Returns a cleanup function that should be called when the screenshot is complete.
+ */
+function prepareCellElementForScreenshot(element: HTMLElement) {
+  element.classList.add("printing-output");
+  acquireBodyPrinting();
+  const originalOverflow = element.style.overflow;
+  element.style.overflow = "auto";
+
+  return () => {
+    element.classList.remove("printing-output");
+    releaseBodyPrinting();
+    element.style.overflow = originalOverflow;
+  };
+}
+
+/**
+ * Capture a cell output as a PNG data URL.
+ */
+export async function getImageDataUrlForCell(
+  cellId: CellId,
+): Promise<string | undefined> {
+  const element = findElementForCell(cellId);
+  if (!element) {
+    return;
+  }
+  const cleanup = prepareCellElementForScreenshot(element);
+
+  try {
+    return await toPng(element);
+  } finally {
+    cleanup();
+  }
+}
+
+/**
+ * Download a cell output as a PNG image file.
+ */
+export async function downloadCellOutputAsImage(
+  cellId: CellId,
   filename: string,
 ) {
+  const element = findElementForCell(cellId);
+  if (!element) {
+    return;
+  }
+
+  await downloadHTMLAsImage({
+    element,
+    filename,
+    prepare: prepareCellElementForScreenshot,
+  });
+}
+
+export async function downloadHTMLAsImage(opts: {
+  element: HTMLElement;
+  filename: string;
+  prepare?: (element: HTMLElement) => () => void;
+}) {
+  const { element, filename, prepare } = opts;
+
   // Capture current scroll position
   const appEl = document.getElementById("App");
   const currentScrollY = appEl?.scrollTop ?? 0;
-  // Add classnames for printing
-  document.body.classList.add("printing");
+
+  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),
+    // add body.printing ourselves
+    document.body.classList.add("printing");
+  }
+
   try {
     // Get screenshot
     const dataUrl = await toPng(element);
@@ -47,8 +146,10 @@ export async function downloadHTMLAsImage(
       variant: "danger",
     });
   } finally {
-    // Remove classnames for printing
-    document.body.classList.remove("printing");
+    cleanup?.();
+    if (document.body.classList.contains("printing")) {
+      document.body.classList.remove("printing");
+    }
     // Restore scroll position
     requestAnimationFrame(() => {
       appEl?.scrollTo(0, currentScrollY);