@marimo-team/frontend 0.19.7-dev30 → 0.19.7-dev33

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/plugins/impl/DataEditorPlugin.tsx

@@ -132,7 +132,7 @@ const LoadingDataEditor = (props: Props) => {
       columnFields={columnFields}
       setColumnFields={setColumnFields}
       editableColumns={props.editableColumns}
-      edits={props.edits.edits} // TODO: This is returning old edits upon refresh
+      edits={props.edits.edits}
       onAddEdits={(edits) => {
         props.onEdits((v) => ({ ...v, edits: [...v.edits, ...edits] }));
       }}

package/src/plugins/impl/data-editor/glide-data-editor.tsx

@@ -14,7 +14,13 @@ import DataEditor, {
   type Rectangle,
 } from "@glideapps/glide-data-grid";
 import { CopyIcon, TrashIcon } from "lucide-react";
-import React, { useCallback, useMemo, useRef, useState } from "react";
+import React, {
+  useCallback,
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+} from "react";
 import useEvent from "react-use-event-hook";
 import type { FieldTypes } from "@/components/data-table/types";
 import {
@@ -40,7 +46,6 @@ import { ErrorBoundary } from "@/components/editor/boundary/ErrorBoundary";
 import { Button } from "@/components/ui/button";
 import { toast } from "@/components/ui/use-toast";
 import type { DataType } from "@/core/kernel/messages";
-import { useOnMount } from "@/hooks/useLifecycle";
 import { useNonce } from "@/hooks/useNonce";
 import { logNever } from "@/utils/assertNever";
 import { Events } from "@/utils/events";
@@ -93,9 +98,47 @@ export const GlideDataEditor = <T,>({
 
   const [columnWidths, setColumnWidths] = useState<Record<string, number>>({});
   const rerender = useNonce();
+  const hasAppliedEdits = useRef(false);
 
-  // Handle initial edits passed in
-  useOnMount(() => {
+  const columns: ModifiedGridColumn[] = useMemo(() => {
+    const columns: ModifiedGridColumn[] = [];
+    for (const [columnName, fieldType] of Object.entries(columnFields)) {
+      const editable =
+        editableColumns === "all" || editableColumns.includes(columnName);
+
+      columns.push({
+        id: columnName,
+        title: columnName,
+        width: columnWidths[columnName], // Enables resizing
+        icon: editable
+          ? getColumnHeaderIcon(fieldType)
+          : GridColumnIcon.ProtectedColumnOverlay,
+        style: "normal",
+        kind: getColumnKind(fieldType),
+        dataType: fieldType,
+        hasMenu: true,
+        themeOverride: editable
+          ? undefined
+          : {
+              bgCell: theme === "light" ? "#F9F9FA" : "#1e1e21",
+            },
+      });
+    }
+
+    return columns;
+  }, [columnFields, columnWidths, editableColumns, theme]);
+
+  // Apply initial edits after data has loaded
+  useEffect(() => {
+    // Don't apply if already applied or data hasn't loaded yet
+    if (hasAppliedEdits.current || data.length === 0) {
+      return;
+    }
+
+    // Mark as applied once data loads - prevents re-applying user edits
+    hasAppliedEdits.current = true;
+
+    // No initial edits to apply
     if (edits.length === 0) {
       return;
     }
@@ -184,35 +227,8 @@ export const GlideDataEditor = <T,>({
 
     // Force re-render to update the total rows
     rerender();
-  });
-
-  const columns: ModifiedGridColumn[] = useMemo(() => {
-    const columns: ModifiedGridColumn[] = [];
-    for (const [columnName, fieldType] of Object.entries(columnFields)) {
-      const editable =
-        editableColumns === "all" || editableColumns.includes(columnName);
-
-      columns.push({
-        id: columnName,
-        title: columnName,
-        width: columnWidths[columnName], // Enables resizing
-        icon: editable
-          ? getColumnHeaderIcon(fieldType)
-          : GridColumnIcon.ProtectedColumnOverlay,
-        style: "normal",
-        kind: getColumnKind(fieldType),
-        dataType: fieldType,
-        hasMenu: true,
-        themeOverride: editable
-          ? undefined
-          : {
-              bgCell: theme === "light" ? "#F9F9FA" : "#1e1e21",
-            },
-      });
-    }
-
-    return columns;
-  }, [columnFields, columnWidths, editableColumns, theme]);
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [data.length]);
 
   const getCellContent = useCallback(
     (cell: Item): GridCell => {

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]);
 

package/src/utils/download.ts

@@ -70,16 +70,15 @@ function releaseBodyPrinting() {
  * Prepare a cell element for screenshot capture.
  *
  * @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.
+ * @param snappy - When true, avoids layout shifts and speeds up the capture.
  * @returns A cleanup function to restore the element's original state
  */
 function prepareCellElementForScreenshot(
   element: HTMLElement,
-  enablePrintMode: boolean,
+  snappy: boolean,
 ) {
   element.classList.add("printing-output");
-  if (enablePrintMode) {
+  if (!snappy) {
     acquireBodyPrinting();
   }
   const originalOverflow = element.style.overflow;
@@ -87,7 +86,7 @@ function prepareCellElementForScreenshot(
 
   return () => {
     element.classList.remove("printing-output");
-    if (enablePrintMode) {
+    if (!snappy) {
       releaseBodyPrinting();
     }
     element.style.overflow = originalOverflow;
@@ -100,13 +99,12 @@ const THRESHOLD_TIME_MS = 500;
  * 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.
+ * @param snappy - When true, uses a faster method to capture the image. Avoids layout shifts.
  * @returns The PNG as a data URL, or undefined if the cell element wasn't found
  */
 export async function getImageDataUrlForCell(
   cellId: CellId,
-  enablePrintMode = true,
+  snappy = false,
 ): Promise<string | undefined> {
   const element = findElementForCell(cellId);
   if (!element) {
@@ -118,11 +116,11 @@ export async function getImageDataUrlForCell(
     return iframeDataUrl;
   }
 
-  const cleanup = prepareCellElementForScreenshot(element, enablePrintMode);
+  const cleanup = prepareCellElementForScreenshot(element, snappy);
 
   try {
     const startTime = Date.now();
-    const dataUrl = await toPng(element);
+    const dataUrl = await toPng(element, undefined, snappy);
     const timeTaken = Date.now() - startTime;
     if (timeTaken > THRESHOLD_TIME_MS) {
       Logger.debug(
@@ -159,7 +157,7 @@ export async function downloadCellOutputAsImage(
   await downloadHTMLAsImage({
     element,
     filename,
-    prepare: () => prepareCellElementForScreenshot(element, true),
+    prepare: () => prepareCellElementForScreenshot(element, false),
   });
 }