@marimo-team/frontend 0.19.7-dev31 → 0.19.7-dev34

package/src/core/export/__tests__/hooks.test.ts

@@ -102,8 +102,8 @@ describe("useEnrichCellOutputs", () => {
 
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
-    const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs(progress);
+    const takeScreenshots = result.current;
+    const output = await takeScreenshots({ progress, snappy: false });
 
     expect(output).toEqual({});
     expect(document.getElementById).not.toHaveBeenCalled();
@@ -134,8 +134,8 @@ describe("useEnrichCellOutputs", () => {
 
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
-    const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs(progress);
+    const takeScreenshots = result.current;
+    const output = await takeScreenshots({ progress, snappy: false });
 
     expect(document.getElementById).toHaveBeenCalledWith(
       CellOutputId.create(cellId),
@@ -152,6 +152,50 @@ describe("useEnrichCellOutputs", () => {
     });
   });
 
+  it("should pass snappy=true to toPng with includeStyleProperties", async () => {
+    const cellId = "cell-1" as CellId;
+    const mockElement = document.createElement("div");
+    const mockDataUrl = "data:image/png;base64,mockImageData";
+
+    // Mock document.getElementById
+    vi.spyOn(document, "getElementById").mockReturnValue(mockElement);
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    setCellsRuntime(
+      createMockCellRuntimes({
+        [cellId]: {
+          output: {
+            channel: "output",
+            mimetype: "text/html",
+            data: "<div>Chart</div>",
+            timestamp: 0,
+          },
+        },
+      }),
+    );
+
+    const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
+
+    const takeScreenshots = result.current;
+    const output = await takeScreenshots({ progress, snappy: true });
+
+    expect(document.getElementById).toHaveBeenCalledWith(
+      CellOutputId.create(cellId),
+    );
+    // When snappy=true, includeStyleProperties should be set
+    expect(toPng).toHaveBeenCalledWith(
+      mockElement,
+      expect.objectContaining({
+        filter: expect.any(Function),
+        onImageErrorHandler: expect.any(Function),
+        includeStyleProperties: expect.any(Array),
+      }),
+    );
+    expect(output).toEqual({
+      [cellId]: ["image/png", mockDataUrl],
+    });
+  });
+
   it("should skip cells where output has not changed", async () => {
     const cellId = "cell-1" as CellId;
     const mockElement = document.createElement("div");
@@ -179,8 +223,8 @@ describe("useEnrichCellOutputs", () => {
     });
 
     // First call - should capture
-    let enrichCellOutputs = result.current;
-    let output = await enrichCellOutputs(progress);
+    let takeScreenshots = result.current;
+    let output = await takeScreenshots({ progress, snappy: false });
     expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl] });
     expect(toPng).toHaveBeenCalledTimes(1);
 
@@ -188,8 +232,8 @@ describe("useEnrichCellOutputs", () => {
     rerender();
 
     // Second call with same output - should not capture again
-    enrichCellOutputs = result.current;
-    output = await enrichCellOutputs(progress);
+    takeScreenshots = result.current;
+    output = await takeScreenshots({ progress, snappy: false });
     expect(output).toEqual({}); // Empty because output hasn't changed
     expect(toPng).toHaveBeenCalledTimes(1); // Still only 1 call
   });
@@ -217,8 +261,8 @@ describe("useEnrichCellOutputs", () => {
 
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
-    const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs(progress);
+    const takeScreenshots = result.current;
+    const output = await takeScreenshots({ progress, snappy: false });
 
     expect(output).toEqual({}); // Failed screenshot should be filtered out
     expect(Logger.error).toHaveBeenCalledWith(
@@ -247,8 +291,8 @@ describe("useEnrichCellOutputs", () => {
 
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
-    const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs(progress);
+    const takeScreenshots = result.current;
+    const output = await takeScreenshots({ progress, snappy: false });
 
     expect(output).toEqual({});
     expect(Logger.error).toHaveBeenCalledWith(
@@ -296,8 +340,8 @@ describe("useEnrichCellOutputs", () => {
 
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
-    const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs(progress);
+    const takeScreenshots = result.current;
+    const output = await takeScreenshots({ progress, snappy: false });
 
     expect(output).toEqual({
       [cell1]: ["image/png", mockDataUrl1],
@@ -342,8 +386,8 @@ describe("useEnrichCellOutputs", () => {
 
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
-    const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs(progress);
+    const takeScreenshots = result.current;
+    const output = await takeScreenshots({ progress, snappy: false });
 
     // Only the successful screenshot should be in the result
     expect(output).toEqual({
@@ -384,14 +428,14 @@ describe("useEnrichCellOutputs", () => {
     });
 
     // First screenshot
-    let enrichCellOutputs = result.current;
-    let output = await enrichCellOutputs(progress);
+    let takeScreenshots = result.current;
+    let output = await takeScreenshots({ progress, snappy: false });
     expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl1] });
 
     // Second call - same output, should not be captured
     rerender();
-    enrichCellOutputs = result.current;
-    output = await enrichCellOutputs(progress);
+    takeScreenshots = result.current;
+    output = await takeScreenshots({ progress, snappy: false });
     expect(output).toEqual({});
 
     // Third call - output changed, should be captured
@@ -409,8 +453,8 @@ describe("useEnrichCellOutputs", () => {
     );
 
     rerender();
-    enrichCellOutputs = result.current;
-    output = await enrichCellOutputs(progress);
+    takeScreenshots = result.current;
+    output = await takeScreenshots({ progress, snappy: false });
     expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl2] });
     expect(toPng).toHaveBeenCalledTimes(2);
   });
@@ -449,8 +493,8 @@ describe("useEnrichCellOutputs", () => {
 
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
-    const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs(progress);
+    const takeScreenshots = result.current;
+    const output = await takeScreenshots({ progress, snappy: false });
 
     // None of these should trigger screenshots
     expect(output).toEqual({});
@@ -474,8 +518,8 @@ describe("useEnrichCellOutputs", () => {
 
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
-    const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs(progress);
+    const takeScreenshots = result.current;
+    const output = await takeScreenshots({ progress, snappy: false });
 
     expect(output).toEqual({});
     expect(document.getElementById).not.toHaveBeenCalled();
@@ -506,8 +550,8 @@ describe("useEnrichCellOutputs", () => {
 
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
-    const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs(progress);
+    const takeScreenshots = result.current;
+    const output = await takeScreenshots({ progress, snappy: false });
 
     // Verify the exact return type structure
     expect(output).toHaveProperty(cellId);
@@ -539,7 +583,6 @@ describe("updateCellOutputsWithScreenshots", () => {
     const updateCellOutputs = vi.fn().mockResolvedValue(null);
 
     await updateCellOutputsWithScreenshots({
-      progress,
       takeScreenshots,
       updateCellOutputs,
     });
@@ -556,7 +599,6 @@ describe("updateCellOutputsWithScreenshots", () => {
     const updateCellOutputs = vi.fn().mockResolvedValue(null);
 
     await updateCellOutputsWithScreenshots({
-      progress,
       takeScreenshots,
       updateCellOutputs,
     });
@@ -583,7 +625,6 @@ describe("updateCellOutputsWithScreenshots", () => {
     const updateCellOutputs = vi.fn().mockResolvedValue(null);
 
     await updateCellOutputsWithScreenshots({
-      progress,
       takeScreenshots,
       updateCellOutputs,
     });
@@ -600,7 +641,6 @@ describe("updateCellOutputsWithScreenshots", () => {
 
     // Should not throw - errors are caught and shown via toast
     await updateCellOutputsWithScreenshots({
-      progress,
       takeScreenshots,
       updateCellOutputs,
     });
@@ -633,7 +673,6 @@ describe("updateCellOutputsWithScreenshots", () => {
 
     // Should not throw - errors are caught and shown via toast
     await updateCellOutputsWithScreenshots({
-      progress,
       takeScreenshots,
       updateCellOutputs,
     });

package/src/core/export/hooks.ts

@@ -64,9 +64,13 @@ export function useAutoExport() {
 
   useInterval(
     async () => {
+      const screenshotFn = () =>
+        takeScreenshots({
+          progress: ProgressState.indeterminate(),
+          snappy: true,
+        });
       await updateCellOutputsWithScreenshots({
-        progress: ProgressState.indeterminate(),
-        takeScreenshots,
+        takeScreenshots: screenshotFn,
         updateCellOutputs,
       });
       await autoExportAsIPYNB({
@@ -107,7 +111,13 @@ export function useEnrichCellOutputs() {
   const [richCellsOutput, setRichCellsOutput] = useAtom(richCellsToOutputAtom);
   const cellRuntimes = useAtomValue(cellsRuntimeAtom);
 
-  return async (progress: ProgressState): Promise<ScreenshotResults> => {
+  return async ({
+    progress,
+    snappy,
+  }: {
+    progress: ProgressState;
+    snappy: boolean;
+  }): Promise<ScreenshotResults> => {
     const trackedCellsOutput: Record<CellId, unknown> = {};
 
     const cellsToCaptureScreenshot: [CellId, unknown][] = [];
@@ -138,7 +148,7 @@ export function useEnrichCellOutputs() {
     const results: ScreenshotResults = {};
     for (const [cellId] of cellsToCaptureScreenshot) {
       try {
-        const dataUrl = await getImageDataUrlForCell(cellId, false);
+        const dataUrl = await getImageDataUrlForCell(cellId, snappy);
         if (!dataUrl) {
           Logger.error(`Failed to capture screenshot for cell ${cellId}`);
           continue;
@@ -159,13 +169,12 @@ export function useEnrichCellOutputs() {
  * Utility function to take screenshots of cells with HTML outputs and update the cell outputs.
  */
 export async function updateCellOutputsWithScreenshots(opts: {
-  progress: ProgressState;
-  takeScreenshots: (progress: ProgressState) => Promise<ScreenshotResults>;
+  takeScreenshots: () => Promise<ScreenshotResults>;
   updateCellOutputs: (request: UpdateCellOutputsRequest) => Promise<null>;
 }) {
-  const { progress, takeScreenshots, updateCellOutputs } = opts;
+  const { takeScreenshots, updateCellOutputs } = opts;
   try {
-    const cellIdsToOutput = await takeScreenshots(progress);
+    const cellIdsToOutput = await takeScreenshots();
     if (Objects.size(cellIdsToOutput) > 0) {
       await updateCellOutputs({ cellIdsToOutput });
     }

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

@@ -166,7 +166,7 @@ describe("getImageDataUrlForCell", () => {
     );
   });
 
-  it("should add printing classes before capture when enablePrintMode is true", async () => {
+  it("should add printing classes before capture when snappy is false", async () => {
     vi.mocked(toPng).mockImplementation(async () => {
       // Check classes are applied during capture
       expect(mockElement.classList.contains("printing-output")).toBe(true);
@@ -175,36 +175,36 @@ describe("getImageDataUrlForCell", () => {
       return mockDataUrl;
     });
 
-    await getImageDataUrlForCell("cell-1" as CellId, true);
+    await getImageDataUrlForCell("cell-1" as CellId, false);
   });
 
-  it("should remove printing classes after capture when enablePrintMode is true", async () => {
+  it("should remove printing classes after capture when snappy is false", async () => {
     vi.mocked(toPng).mockResolvedValue(mockDataUrl);
 
-    await getImageDataUrlForCell("cell-1" as CellId, true);
+    await getImageDataUrlForCell("cell-1" as CellId, false);
 
     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 () => {
+  it("should add printing-output but NOT body.printing when snappy is true", 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
+      // but body.printing should NOT be added when snappy mode is on
       expect(document.body.classList.contains("printing")).toBe(false);
       expect(mockElement.style.overflow).toBe("auto");
       return mockDataUrl;
     });
 
-    await getImageDataUrlForCell("cell-1" as CellId, false);
+    await getImageDataUrlForCell("cell-1" as CellId, true);
   });
 
-  it("should cleanup printing-output when enablePrintMode is false", async () => {
+  it("should cleanup printing-output when snappy is true", async () => {
     mockElement.style.overflow = "hidden";
     vi.mocked(toPng).mockResolvedValue(mockDataUrl);
 
-    await getImageDataUrlForCell("cell-1" as CellId, false);
+    await getImageDataUrlForCell("cell-1" as CellId, true);
 
     expect(mockElement.classList.contains("printing-output")).toBe(false);
     expect(document.body.classList.contains("printing")).toBe(false);
@@ -239,7 +239,7 @@ describe("getImageDataUrlForCell", () => {
     expect(mockElement.style.overflow).toBe("scroll");
   });
 
-  it("should maintain body.printing during concurrent captures when enablePrintMode is true", async () => {
+  it("should maintain body.printing during concurrent captures when snappy is false", async () => {
     // Create a second element
     const mockElement2 = document.createElement("div");
     mockElement2.id = CellOutputId.create("cell-2" as CellId);
@@ -273,9 +273,9 @@ describe("getImageDataUrlForCell", () => {
       return mockDataUrl;
     });
 
-    // Start both captures concurrently with enablePrintMode = true
-    const capture1 = getImageDataUrlForCell("cell-1" as CellId, true);
-    const capture2 = getImageDataUrlForCell("cell-2" as CellId, true);
+    // Start both captures concurrently with snappy = false (body.printing should be added)
+    const capture1 = getImageDataUrlForCell("cell-1" as CellId, false);
+    const capture2 = getImageDataUrlForCell("cell-2" as CellId, false);
 
     // Let second capture complete first
     resolveSecond!();
@@ -297,21 +297,21 @@ describe("getImageDataUrlForCell", () => {
     mockElement2.remove();
   });
 
-  it("should not interfere with body.printing during concurrent captures when enablePrintMode is false", async () => {
+  it("should not interfere with body.printing during concurrent captures when snappy is true", 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
+      // body.printing should never be added when snappy is true
       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);
+    // Start both captures concurrently with snappy = true (body.printing should NOT be added)
+    const capture1 = getImageDataUrlForCell("cell-1" as CellId, true);
+    const capture2 = getImageDataUrlForCell("cell-2" as CellId, true);
 
     await Promise.all([capture1, capture2]);