Merge pull request #2173 from mito-ds/feat/chart-wizard-export

Chart Wizard Export

Author: aarondr77

mito-ai/src/Extensions/ChartWizard/ChartWizardPlugin.tsx

@@ -15,7 +15,7 @@ import { Widget } from '@lumino/widgets';
 import { ChartWizardWidget } from './ChartWizardWidget';
 import { COMMAND_MITO_AI_OPEN_CHART_WIZARD } from '../../commands';
 import TextAndIconButton from '../../components/TextAndIconButton';
-import MagicWand from '../../icons/MagicWand'
+import MagicWand from '../../icons/MagicWand';
 import { logEvent } from '../../restAPI/RestAPI';
 import { setActiveCellByIDInNotebookPanel } from '../../utils/notebook';
 import '../../../style/ChartWizardPlugin.css'
@@ -34,18 +34,16 @@ interface ChartWizardButtonProps {
 
 const ChartWizardButton: React.FC<ChartWizardButtonProps> = ({ onButtonClick }) => {
     return (
-        <>
-            <TextAndIconButton
-                icon={MagicWand}
-                text="Chart Wizard"
-                title="Chart Wizard"
-                onClick={onButtonClick}
-                variant='purple'
-                width='fit-contents'
-                iconPosition='left'
-            />
-        </>
-    )
+        <TextAndIconButton
+            icon={MagicWand}
+            text="Chart Wizard"
+            title="Chart Wizard"
+            onClick={onButtonClick}
+            variant="purple"
+            width="fit-contents"
+            iconPosition="left"
+        />
+    );
 };
 
 const ChartWizardPlugin: JupyterFrontEndPlugin<void> = {
@@ -202,7 +200,7 @@ class AugmentedImageRenderer extends Widget implements IRenderMime.IRenderer {
         super.dispose();
     }
 
-    /* 
+    /*
         Handle the Chart Wizard button click.
         Extracts chart data and source code, then opens the Chart Wizard panel.
     */

mito-ai/src/Extensions/ChartWizard/ChartWizardWidget.tsx

@@ -7,9 +7,11 @@ import React, { useState, useEffect, useCallback, useMemo, useRef } from 'react'
 import { ReactWidget, Notification } from '@jupyterlab/apputils';
 import { ChartWizardData } from './ChartWizardPlugin';
 import { updateChartConfig, ChartConfigVariable } from './utils/parser';
+import { exportChartImage, type ExportImageFormat } from './utils/chartExport';
 import { convertChartCode, logEvent } from '../../restAPI/RestAPI';
 import { removeMarkdownCodeFormatting } from '../../utils/strings';
 import LoadingDots from '../../components/LoadingDots';
+import ToggleButton from '../../components/ToggleButton';
 import AddFieldButton from './AddFieldButton';
 import {
     BooleanInputRow,
@@ -45,6 +47,7 @@ const ChartWizardContent: React.FC<ChartWizardContentProps> = ({ chartData }) =>
     const widgetRef = useRef<HTMLDivElement>(null);
     const [overlayHeight, setOverlayHeight] = useState<number>(0);
     const [isActiveCellMismatch, setIsActiveCellMismatch] = useState(false);
+    const [exportFormat, setExportFormat] = useState<ExportImageFormat>('png');
 
     // Reset currentSourceCode when switching to a different chart
     useEffect(() => {
@@ -214,6 +217,19 @@ const ChartWizardContent: React.FC<ChartWizardContentProps> = ({ chartData }) =>
         updateNotebookCell(updatedCode);
     }, [updateNotebookCell]);
 
+    /**
+     * Exports the chart image to disk via the chartExport utility; shows a notification on error.
+     */
+    const handleExportChart = useCallback(async (): Promise<void> => {
+        void logEvent('chart_wizard_export_clicked');
+        if (!chartData) return;
+
+        const result = await exportChartImage(chartData, exportFormat);
+        if (!result.success) {
+            Notification.emit(result.error, 'error', { autoClose: 5000 });
+        }
+    }, [chartData, exportFormat]);
+
     /**
      * Renders the appropriate input field component based on variable type.
      */
@@ -319,6 +335,25 @@ const ChartWizardContent: React.FC<ChartWizardContentProps> = ({ chartData }) =>
                         clearPendingUpdate={clearPendingUpdate}
                         onLoadingStateChange={setIsAddingField}
                     />
+                    <div className="chart-wizard-export-section">
+                        <h3 className="chart-wizard-section-heading">Export</h3>
+                        <div className="chart-wizard-export-format-row">
+                            <ToggleButton
+                                leftText="PNG"
+                                rightText="JPG"
+                                isLeftSelected={exportFormat === 'png'}
+                                onChange={(isPng) => setExportFormat(isPng ? 'png' : 'jpeg')}
+                            />
+                        </div>
+                        <button
+                            className="button-base button-purple add-field-button"
+                            type="button"
+                            title="Save chart image to file"
+                            onClick={handleExportChart}
+                        >
+                            Export image
+                        </button>
+                    </div>
                 </div>
             ) : (
                 <div className="chart-wizard-no-config">

mito-ai/src/Extensions/ChartWizard/utils/chartExport.ts

@@ -0,0 +1,45 @@
+/*
+ * Copyright (c) Saga Inc.
+ * Distributed under the terms of the GNU Affero General Public License v3.0 License.
+ */
+
+import { ChartWizardData } from '../ChartWizardPlugin';
+import { findChartImageDataUrl } from './imageFinder';
+import { saveWithFilePicker, isFileSystemAccessAvailable } from './fileSaver';
+import { downloadImage } from './download';
+import { ExportImageFormat } from './types';
+
+export type ExportChartResult = { success: true } | { success: false; error: string };
+
+export type { ExportImageFormat };
+
+/**
+ * Exports the chart image to the user's disk. Uses File System Access API when available
+ * so the user can choose the save location; otherwise triggers a download.
+ *
+ * @param chartData - Chart wizard data identifying the notebook panel and cell
+ * @param format - Export format: 'png' or 'jpeg'
+ * @returns Result indicating success or an error message for the UI to display
+ */
+export async function exportChartImage(
+    chartData: ChartWizardData,
+    format: ExportImageFormat = 'png'
+): Promise<ExportChartResult> {
+    const found = findChartImageDataUrl(chartData);
+    if (!found.ok) return { success: false, error: found.error };
+
+    if (isFileSystemAccessAvailable()) {
+        try {
+            await saveWithFilePicker(found.dataUrl, format);
+        } catch (err) {
+            if ((err as { name?: string }).name === 'AbortError') {
+                return { success: true };
+            }
+            await downloadImage(found.dataUrl, format);
+        }
+    } else {
+        await downloadImage(found.dataUrl, format);
+    }
+
+    return { success: true };
+}

mito-ai/src/Extensions/ChartWizard/utils/download.ts

@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) Saga Inc.
+ * Distributed under the terms of the GNU Affero General Public License v3.0 License.
+ */
+
+import { ExportImageFormat } from './types';
+import { dataUrlToBlob } from './imageConverter';
+
+const SUGGESTED_NAMES: Record<ExportImageFormat, string> = {
+    png: 'chart.png',
+    jpeg: 'chart.jpg'
+};
+
+/**
+ * Triggers a browser download of a URL.
+ */
+function triggerDownload(url: string, filename: string): void {
+    const a = document.createElement('a');
+    a.href = url;
+    a.download = filename;
+    a.click();
+}
+
+/**
+ * Downloads an image using the browser's fallback download mechanism.
+ */
+export async function downloadImage(
+    dataUrl: string,
+    format: ExportImageFormat
+): Promise<void> {
+    if (format === 'jpeg') {
+        const blob = await dataUrlToBlob(dataUrl, format);
+        const url = URL.createObjectURL(blob);
+        triggerDownload(url, SUGGESTED_NAMES.jpeg);
+        URL.revokeObjectURL(url);
+    } else {
+        triggerDownload(dataUrl, SUGGESTED_NAMES.png);
+    }
+}

mito-ai/src/Extensions/ChartWizard/utils/fileSaver.ts

@@ -0,0 +1,70 @@
+/*
+ * Copyright (c) Saga Inc.
+ * Distributed under the terms of the GNU Affero General Public License v3.0 License.
+ */
+
+import { ExportImageFormat } from './types';
+import { dataUrlToBlob } from './imageConverter';
+
+const SUGGESTED_NAMES: Record<ExportImageFormat, string> = {
+    png: 'chart.png',
+    jpeg: 'chart.jpg'
+};
+
+const FILE_PICKER_TYPES: Record<
+    ExportImageFormat,
+    Array<{ description: string; accept: Record<string, string[]> }>
+> = {
+    png: [{ description: 'PNG Image', accept: { 'image/png': ['.png'] } }],
+    jpeg: [{ description: 'JPEG Image', accept: { 'image/jpeg': ['.jpg', '.jpeg'] } }]
+};
+
+/**
+ * Checks if the File System Access API is available.
+ */
+export function isFileSystemAccessAvailable(): boolean {
+    return (
+        'showSaveFilePicker' in window &&
+        typeof (window as Window & { showSaveFilePicker?: unknown }).showSaveFilePicker ===
+            'function'
+    );
+}
+
+/**
+ * Saves a blob to disk using the File System Access API.
+ */
+async function writeBlobToFile(blob: Blob, handle: FileSystemFileHandle): Promise<void> {
+    const writable = await (handle as FileSystemFileHandle & {
+        createWritable(): Promise<{ write(data: Blob): Promise<void>; close(): Promise<void> }>;
+    }).createWritable();
+    try {
+        await writable.write(blob);
+    } finally {
+        await writable.close();
+    }
+}
+
+/**
+ * Saves an image to disk using the File System Access API.
+ */
+export async function saveWithFilePicker(
+    dataUrl: string,
+    format: ExportImageFormat
+): Promise<void> {
+    const handle = await (window as Window & {
+        showSaveFilePicker?: (options: {
+            suggestedName?: string;
+            types?: Array<{
+                description: string;
+                accept: Record<string, string[]>;
+            }>;
+        }) => Promise<FileSystemFileHandle>;
+    }).showSaveFilePicker?.({
+        suggestedName: SUGGESTED_NAMES[format],
+        types: FILE_PICKER_TYPES[format]
+    });
+    if (!handle) return;
+
+    const blob = await dataUrlToBlob(dataUrl, format);
+    await writeBlobToFile(blob, handle);
+}

mito-ai/src/Extensions/ChartWizard/utils/imageConverter.ts

@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) Saga Inc.
+ * Distributed under the terms of the GNU Affero General Public License v3.0 License.
+ */
+
+import { ExportImageFormat } from './types';
+
+const JPEG_QUALITY = 1.0;
+
+/**
+ * Converts a data URL image to a JPEG blob.
+ */
+export function dataUrlToJpegBlob(dataUrl: string): Promise<Blob> {
+    return new Promise((resolve, reject) => {
+        const img = new Image();
+        img.crossOrigin = 'anonymous';
+        img.onload = (): void => {
+            const canvas = document.createElement('canvas');
+            canvas.width = img.naturalWidth;
+            canvas.height = img.naturalHeight;
+            const ctx = canvas.getContext('2d');
+            if (!ctx) {
+                reject(new Error('Could not get canvas context'));
+                return;
+            }
+            ctx.drawImage(img, 0, 0);
+            canvas.toBlob(
+                (blob) => (blob ? resolve(blob) : reject(new Error('toBlob failed'))),
+                'image/jpeg',
+                JPEG_QUALITY
+            );
+        };
+        img.onerror = (): void => reject(new Error('Failed to load image'));
+        img.src = dataUrl;
+    });
+}
+
+/**
+ * Converts a data URL to a blob in the specified format.
+ */
+export async function dataUrlToBlob(
+    dataUrl: string,
+    format: ExportImageFormat
+): Promise<Blob> {
+    if (format === 'jpeg') {
+        return dataUrlToJpegBlob(dataUrl);
+    }
+    return fetch(dataUrl).then((r) => r.blob());
+}

mito-ai/src/Extensions/ChartWizard/utils/imageFinder.ts

@@ -0,0 +1,43 @@
+/*
+ * Copyright (c) Saga Inc.
+ * Distributed under the terms of the GNU Affero General Public License v3.0 License.
+ */
+
+import { CodeCell } from '@jupyterlab/cells';
+import { ChartWizardData } from '../ChartWizardPlugin';
+
+export type FindImageResult =
+    | { ok: true; dataUrl: string }
+    | { ok: false; error: string };
+
+/**
+ * Finds the chart image data URL from the notebook cell output.
+ */
+export function findChartImageDataUrl(chartData: ChartWizardData): FindImageResult {
+    const notebookPanel = chartData.notebookTracker.find(
+        (panel) => panel.id === chartData.notebookPanelId
+    );
+    if (!notebookPanel) {
+        return { ok: false, error: 'Could not find the notebook.' };
+    }
+
+    const cellWidget = notebookPanel.content.widgets.find(
+        (cell) => cell.model.id === chartData.cellId
+    );
+    if (!(cellWidget instanceof CodeCell)) {
+        return { ok: false, error: 'Could not find the chart cell.' };
+    }
+
+    const outputNode = cellWidget.outputArea.node;
+    const img = outputNode.querySelector(
+        '.jp-RenderedImage img[src^="data:image"]'
+    ) as HTMLImageElement | null;
+
+    if (!img || !img.src || !img.src.startsWith('data:image')) {
+        return {
+            ok: false,
+            error: 'No chart image found. Re-run the chart cell and try again.'
+        };
+    }
+    return { ok: true, dataUrl: img.src };
+}

mito-ai/src/Extensions/ChartWizard/utils/types.ts

@@ -0,0 +1,6 @@
+/*
+ * Copyright (c) Saga Inc.
+ * Distributed under the terms of the GNU Affero General Public License v3.0 License.
+ */
+
+export type ExportImageFormat = 'png' | 'jpeg';

mito-ai/style/AddFieldButton.css

@@ -5,8 +5,6 @@
 
 .add-field-container {
     margin-top: 16px;
-    padding-top: 16px;
-    border-top: 1px solid #e0e0e0;
 }
 
 .add-field-button {