Skip to content

Commit 9af83cc

Browse files
Light2DarkLight2Dark
authored
add snappy parameter to speed up auto-export (#8032)
## 📝 Summary <!-- Provide a concise summary of what this pull request is addressing. If this PR closes any issues, list them here by number (e.g., Closes #123). --> Adds a new 'snappy' parameter to speed up auto-export to ipynb. This makes it so that taking screenshots will only look for necessary css styles (which offers a lot of speedup) ## 🔍 Description of Changes <!-- Detail the specific changes made in this pull request. Explain the problem addressed and how it was resolved. If applicable, provide before and after comparisons, screenshots, or any relevant details to help reviewers understand the changes easily. --> ## 📋 Checklist - [x] I have read the [contributor guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md). - [x] For large changes, or changes that affect the public API: this change was discussed or approved through an issue, on [Discord](https://marimo.io/discord?ref=pr), or the community [discussions](https://github.com/marimo-team/marimo/discussions) (Please provide a link if applicable). - [x] Tests have been added for the changes made. - [ ] Documentation has been updated where applicable, including docstrings for API changes. - [x] Pull request title is a good summary of the changes - it will be used in the [release notes](https://github.com/marimo-team/marimo/releases).
1 parent d72a1a3 commit 9af83cc

File tree

6 files changed

+257
-72
lines changed

6 files changed

+257
-72
lines changed

‎frontend/src/components/editor/actions/useNotebookActions.tsx‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -241,8 +241,8 @@ export function useNotebookActions() {
241241

242242
const downloadPDF = async (progress: ProgressState) => {
243243
await updateCellOutputsWithScreenshots({
244-
progress,
245-
takeScreenshots,
244+
takeScreenshots: () =>
245+
takeScreenshots({ progress, snappy: false }),
246246
updateCellOutputs,
247247
});
248248
await downloadAsPDF({

‎frontend/src/core/export/__tests__/hooks.test.ts‎

Lines changed: 72 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -102,8 +102,8 @@ describe("useEnrichCellOutputs", () => {
102102

103103
const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
104104

105-
const enrichCellOutputs = result.current;
106-
const output = await enrichCellOutputs(progress);
105+
const takeScreenshots = result.current;
106+
const output = await takeScreenshots({ progress, snappy: false });
107107

108108
expect(output).toEqual({});
109109
expect(document.getElementById).not.toHaveBeenCalled();
@@ -134,8 +134,8 @@ describe("useEnrichCellOutputs", () => {
134134

135135
const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
136136

137-
const enrichCellOutputs = result.current;
138-
const output = await enrichCellOutputs(progress);
137+
const takeScreenshots = result.current;
138+
const output = await takeScreenshots({ progress, snappy: false });
139139

140140
expect(document.getElementById).toHaveBeenCalledWith(
141141
CellOutputId.create(cellId),
@@ -152,6 +152,50 @@ describe("useEnrichCellOutputs", () => {
152152
});
153153
});
154154

155+
it("should pass snappy=true to toPng with includeStyleProperties", async () => {
156+
const cellId = "cell-1" as CellId;
157+
const mockElement = document.createElement("div");
158+
const mockDataUrl = "data:image/png;base64,mockImageData";
159+
160+
// Mock document.getElementById
161+
vi.spyOn(document, "getElementById").mockReturnValue(mockElement);
162+
vi.mocked(toPng).mockResolvedValue(mockDataUrl);
163+
164+
setCellsRuntime(
165+
createMockCellRuntimes({
166+
[cellId]: {
167+
output: {
168+
channel: "output",
169+
mimetype: "text/html",
170+
data: "<div>Chart</div>",
171+
timestamp: 0,
172+
},
173+
},
174+
}),
175+
);
176+
177+
const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
178+
179+
const takeScreenshots = result.current;
180+
const output = await takeScreenshots({ progress, snappy: true });
181+
182+
expect(document.getElementById).toHaveBeenCalledWith(
183+
CellOutputId.create(cellId),
184+
);
185+
// When snappy=true, includeStyleProperties should be set
186+
expect(toPng).toHaveBeenCalledWith(
187+
mockElement,
188+
expect.objectContaining({
189+
filter: expect.any(Function),
190+
onImageErrorHandler: expect.any(Function),
191+
includeStyleProperties: expect.any(Array),
192+
}),
193+
);
194+
expect(output).toEqual({
195+
[cellId]: ["image/png", mockDataUrl],
196+
});
197+
});
198+
155199
it("should skip cells where output has not changed", async () => {
156200
const cellId = "cell-1" as CellId;
157201
const mockElement = document.createElement("div");
@@ -179,17 +223,17 @@ describe("useEnrichCellOutputs", () => {
179223
});
180224

181225
// First call - should capture
182-
let enrichCellOutputs = result.current;
183-
let output = await enrichCellOutputs(progress);
226+
let takeScreenshots = result.current;
227+
let output = await takeScreenshots({ progress, snappy: false });
184228
expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl] });
185229
expect(toPng).toHaveBeenCalledTimes(1);
186230

187231
// Rerender to get updated atom state
188232
rerender();
189233

190234
// Second call with same output - should not capture again
191-
enrichCellOutputs = result.current;
192-
output = await enrichCellOutputs(progress);
235+
takeScreenshots = result.current;
236+
output = await takeScreenshots({ progress, snappy: false });
193237
expect(output).toEqual({}); // Empty because output hasn't changed
194238
expect(toPng).toHaveBeenCalledTimes(1); // Still only 1 call
195239
});
@@ -217,8 +261,8 @@ describe("useEnrichCellOutputs", () => {
217261

218262
const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
219263

220-
const enrichCellOutputs = result.current;
221-
const output = await enrichCellOutputs(progress);
264+
const takeScreenshots = result.current;
265+
const output = await takeScreenshots({ progress, snappy: false });
222266

223267
expect(output).toEqual({}); // Failed screenshot should be filtered out
224268
expect(Logger.error).toHaveBeenCalledWith(
@@ -247,8 +291,8 @@ describe("useEnrichCellOutputs", () => {
247291

248292
const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
249293

250-
const enrichCellOutputs = result.current;
251-
const output = await enrichCellOutputs(progress);
294+
const takeScreenshots = result.current;
295+
const output = await takeScreenshots({ progress, snappy: false });
252296

253297
expect(output).toEqual({});
254298
expect(Logger.error).toHaveBeenCalledWith(
@@ -296,8 +340,8 @@ describe("useEnrichCellOutputs", () => {
296340

297341
const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
298342

299-
const enrichCellOutputs = result.current;
300-
const output = await enrichCellOutputs(progress);
343+
const takeScreenshots = result.current;
344+
const output = await takeScreenshots({ progress, snappy: false });
301345

302346
expect(output).toEqual({
303347
[cell1]: ["image/png", mockDataUrl1],
@@ -342,8 +386,8 @@ describe("useEnrichCellOutputs", () => {
342386

343387
const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
344388

345-
const enrichCellOutputs = result.current;
346-
const output = await enrichCellOutputs(progress);
389+
const takeScreenshots = result.current;
390+
const output = await takeScreenshots({ progress, snappy: false });
347391

348392
// Only the successful screenshot should be in the result
349393
expect(output).toEqual({
@@ -384,14 +428,14 @@ describe("useEnrichCellOutputs", () => {
384428
});
385429

386430
// First screenshot
387-
let enrichCellOutputs = result.current;
388-
let output = await enrichCellOutputs(progress);
431+
let takeScreenshots = result.current;
432+
let output = await takeScreenshots({ progress, snappy: false });
389433
expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl1] });
390434

391435
// Second call - same output, should not be captured
392436
rerender();
393-
enrichCellOutputs = result.current;
394-
output = await enrichCellOutputs(progress);
437+
takeScreenshots = result.current;
438+
output = await takeScreenshots({ progress, snappy: false });
395439
expect(output).toEqual({});
396440

397441
// Third call - output changed, should be captured
@@ -409,8 +453,8 @@ describe("useEnrichCellOutputs", () => {
409453
);
410454

411455
rerender();
412-
enrichCellOutputs = result.current;
413-
output = await enrichCellOutputs(progress);
456+
takeScreenshots = result.current;
457+
output = await takeScreenshots({ progress, snappy: false });
414458
expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl2] });
415459
expect(toPng).toHaveBeenCalledTimes(2);
416460
});
@@ -449,8 +493,8 @@ describe("useEnrichCellOutputs", () => {
449493

450494
const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
451495

452-
const enrichCellOutputs = result.current;
453-
const output = await enrichCellOutputs(progress);
496+
const takeScreenshots = result.current;
497+
const output = await takeScreenshots({ progress, snappy: false });
454498

455499
// None of these should trigger screenshots
456500
expect(output).toEqual({});
@@ -474,8 +518,8 @@ describe("useEnrichCellOutputs", () => {
474518

475519
const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
476520

477-
const enrichCellOutputs = result.current;
478-
const output = await enrichCellOutputs(progress);
521+
const takeScreenshots = result.current;
522+
const output = await takeScreenshots({ progress, snappy: false });
479523

480524
expect(output).toEqual({});
481525
expect(document.getElementById).not.toHaveBeenCalled();
@@ -506,8 +550,8 @@ describe("useEnrichCellOutputs", () => {
506550

507551
const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
508552

509-
const enrichCellOutputs = result.current;
510-
const output = await enrichCellOutputs(progress);
553+
const takeScreenshots = result.current;
554+
const output = await takeScreenshots({ progress, snappy: false });
511555

512556
// Verify the exact return type structure
513557
expect(output).toHaveProperty(cellId);
@@ -539,7 +583,6 @@ describe("updateCellOutputsWithScreenshots", () => {
539583
const updateCellOutputs = vi.fn().mockResolvedValue(null);
540584

541585
await updateCellOutputsWithScreenshots({
542-
progress,
543586
takeScreenshots,
544587
updateCellOutputs,
545588
});
@@ -556,7 +599,6 @@ describe("updateCellOutputsWithScreenshots", () => {
556599
const updateCellOutputs = vi.fn().mockResolvedValue(null);
557600

558601
await updateCellOutputsWithScreenshots({
559-
progress,
560602
takeScreenshots,
561603
updateCellOutputs,
562604
});
@@ -583,7 +625,6 @@ describe("updateCellOutputsWithScreenshots", () => {
583625
const updateCellOutputs = vi.fn().mockResolvedValue(null);
584626

585627
await updateCellOutputsWithScreenshots({
586-
progress,
587628
takeScreenshots,
588629
updateCellOutputs,
589630
});
@@ -600,7 +641,6 @@ describe("updateCellOutputsWithScreenshots", () => {
600641

601642
// Should not throw - errors are caught and shown via toast
602643
await updateCellOutputsWithScreenshots({
603-
progress,
604644
takeScreenshots,
605645
updateCellOutputs,
606646
});
@@ -633,7 +673,6 @@ describe("updateCellOutputsWithScreenshots", () => {
633673

634674
// Should not throw - errors are caught and shown via toast
635675
await updateCellOutputsWithScreenshots({
636-
progress,
637676
takeScreenshots,
638677
updateCellOutputs,
639678
});

‎frontend/src/core/export/hooks.ts‎

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -64,9 +64,13 @@ export function useAutoExport() {
6464

6565
useInterval(
6666
async () => {
67+
const screenshotFn = () =>
68+
takeScreenshots({
69+
progress: ProgressState.indeterminate(),
70+
snappy: true,
71+
});
6772
await updateCellOutputsWithScreenshots({
68-
progress: ProgressState.indeterminate(),
69-
takeScreenshots,
73+
takeScreenshots: screenshotFn,
7074
updateCellOutputs,
7175
});
7276
await autoExportAsIPYNB({
@@ -107,7 +111,13 @@ export function useEnrichCellOutputs() {
107111
const [richCellsOutput, setRichCellsOutput] = useAtom(richCellsToOutputAtom);
108112
const cellRuntimes = useAtomValue(cellsRuntimeAtom);
109113

110-
return async (progress: ProgressState): Promise<ScreenshotResults> => {
114+
return async ({
115+
progress,
116+
snappy,
117+
}: {
118+
progress: ProgressState;
119+
snappy: boolean;
120+
}): Promise<ScreenshotResults> => {
111121
const trackedCellsOutput: Record<CellId, unknown> = {};
112122

113123
const cellsToCaptureScreenshot: [CellId, unknown][] = [];
@@ -138,7 +148,7 @@ export function useEnrichCellOutputs() {
138148
const results: ScreenshotResults = {};
139149
for (const [cellId] of cellsToCaptureScreenshot) {
140150
try {
141-
const dataUrl = await getImageDataUrlForCell(cellId, false);
151+
const dataUrl = await getImageDataUrlForCell(cellId, snappy);
142152
if (!dataUrl) {
143153
Logger.error(`Failed to capture screenshot for cell ${cellId}`);
144154
continue;
@@ -159,13 +169,12 @@ export function useEnrichCellOutputs() {
159169
* Utility function to take screenshots of cells with HTML outputs and update the cell outputs.
160170
*/
161171
export async function updateCellOutputsWithScreenshots(opts: {
162-
progress: ProgressState;
163-
takeScreenshots: (progress: ProgressState) => Promise<ScreenshotResults>;
172+
takeScreenshots: () => Promise<ScreenshotResults>;
164173
updateCellOutputs: (request: UpdateCellOutputsRequest) => Promise<null>;
165174
}) {
166-
const { progress, takeScreenshots, updateCellOutputs } = opts;
175+
const { takeScreenshots, updateCellOutputs } = opts;
167176
try {
168-
const cellIdsToOutput = await takeScreenshots(progress);
177+
const cellIdsToOutput = await takeScreenshots();
169178
if (Objects.size(cellIdsToOutput) > 0) {
170179
await updateCellOutputs({ cellIdsToOutput });
171180
}

0 commit comments

Comments
 (0)