From 5da8ac480c33269ffef9818cee4ebe67dc4f25c0 Mon Sep 17 00:00:00 2001 From: Christoffer Artmann Date: Wed, 15 Oct 2025 14:21:40 +0200 Subject: [PATCH 1/2] feat: Copy and export data frames. --- package-lock.json | 46 +- package.json | 2 + .../dataframe/dataframeController.ts | 275 +++++-- .../dataframeController.unit.test.ts | 742 ++++++++++++++++++ .../dataframe-renderer/DataframeRenderer.tsx | 96 ++- 5 files changed, 1071 insertions(+), 90 deletions(-) create mode 100644 src/webviews/extension-side/dataframe/dataframeController.unit.test.ts diff --git a/package-lock.json b/package-lock.json index 28ab4f22db..6a048d72f7 100644 --- a/package-lock.json +++ b/package-lock.json @@ -26,6 +26,7 @@ "ansi-to-html": "^0.6.7", "bootstrap": "^5.0.0", "bootstrap-less": "^3.3.8", + "clsx": "^2.1.1", "cross-fetch": "^3.1.5", "encoding": "^0.1.13", "fast-deep-equal": "^2.0.1", @@ -64,6 +65,7 @@ "strip-comments": "^2.0.1", "styled-components": "^5.2.1", "svg-to-pdfkit": "^0.1.8", + "tailwind-merge": "^3.3.1", "tcp-port-used": "^1.0.1", "tmp": "^0.2.4", "url-parse": "^1.5.10", @@ -6032,9 +6034,10 @@ } }, "node_modules/clsx": { - "version": "1.1.1", - "resolved": "https://registry.npmjs.org/clsx/-/clsx-1.1.1.tgz", - "integrity": "sha512-6/bPho624p3S2pMyvP5kKBPXnI3ufHLObBFCfgx+LkeR5lg2XYy2hqZqUf45ypD8COn2bhgGJSUE+l5dhNBieA==", + "version": "2.1.1", + "resolved": "https://registry.npmjs.org/clsx/-/clsx-2.1.1.tgz", + "integrity": "sha512-eYm0QWBtUrBWZWG0d386OGAw16Z995PiOVo2B7bjWSbHedGl5e0ZWaq65kOGgUSNesEIDkB9ISbTg/JK9dhCZA==", + "license": "MIT", "engines": { "node": ">=6" } @@ -14839,6 +14842,15 @@ "react-dom": "^16.3.0 || ^17.0.0 || ^18.0.0 || ^19.0.0" } }, + "node_modules/react-virtualized/node_modules/clsx": { + "version": "1.2.1", + "resolved": "https://registry.npmjs.org/clsx/-/clsx-1.2.1.tgz", + "integrity": "sha512-EcR6r5a8bj6pu3ycsa/E/cKVGuTgZJZdsyUYHOksG/UHIiKfjxzRxYJpyVBwYaQeOvghal9fcc4PidlgzugAQg==", + "license": "MIT", + "engines": { + "node": ">=6" + } + }, "node_modules/readable-stream": { "version": "3.6.2", "resolved": "https://registry.npmjs.org/readable-stream/-/readable-stream-3.6.2.tgz", @@ -16296,6 +16308,16 @@ "integrity": "sha512-AEYxH93jGFPn/a2iVAwW87VuUIkR1FVUKB77NwMF7nBTDkDrrT/Hpt/IrCJ0QXhW27jTBDcf5ZY7w6RiqTMw2Q==", "dev": true }, + "node_modules/tailwind-merge": { + "version": "3.3.1", + "resolved": "https://registry.npmjs.org/tailwind-merge/-/tailwind-merge-3.3.1.tgz", + "integrity": "sha512-gBXpgUm/3rp1lMZZrM/w7D8GKqshif0zAymAhbCyIt8KMe+0v9DQ7cdYLR4FHH/cKpdTXb+A/tKKU3eolfsI+g==", + "license": "MIT", + "funding": { + "type": "github", + "url": "https://github.com/sponsors/dcastil" + } + }, "node_modules/tailwindcss": { "version": "4.1.14", "resolved": "https://registry.npmjs.org/tailwindcss/-/tailwindcss-4.1.14.tgz", @@ -22329,9 +22351,9 @@ } }, "clsx": { - "version": "1.1.1", - "resolved": "https://registry.npmjs.org/clsx/-/clsx-1.1.1.tgz", - "integrity": "sha512-6/bPho624p3S2pMyvP5kKBPXnI3ufHLObBFCfgx+LkeR5lg2XYy2hqZqUf45ypD8COn2bhgGJSUE+l5dhNBieA==" + "version": "2.1.1", + "resolved": "https://registry.npmjs.org/clsx/-/clsx-2.1.1.tgz", + "integrity": "sha512-eYm0QWBtUrBWZWG0d386OGAw16Z995PiOVo2B7bjWSbHedGl5e0ZWaq65kOGgUSNesEIDkB9ISbTg/JK9dhCZA==" }, "cmake-ts": { "version": "1.0.2", @@ -28853,6 +28875,13 @@ "loose-envify": "^1.4.0", "prop-types": "^15.7.2", "react-lifecycles-compat": "^3.0.4" + }, + "dependencies": { + "clsx": { + "version": "1.2.1", + "resolved": "https://registry.npmjs.org/clsx/-/clsx-1.2.1.tgz", + "integrity": "sha512-EcR6r5a8bj6pu3ycsa/E/cKVGuTgZJZdsyUYHOksG/UHIiKfjxzRxYJpyVBwYaQeOvghal9fcc4PidlgzugAQg==" + } } }, "readable-stream": { @@ -29986,6 +30015,11 @@ } } }, + "tailwind-merge": { + "version": "3.3.1", + "resolved": "https://registry.npmjs.org/tailwind-merge/-/tailwind-merge-3.3.1.tgz", + "integrity": "sha512-gBXpgUm/3rp1lMZZrM/w7D8GKqshif0zAymAhbCyIt8KMe+0v9DQ7cdYLR4FHH/cKpdTXb+A/tKKU3eolfsI+g==" + }, "tailwindcss": { "version": "4.1.14", "resolved": "https://registry.npmjs.org/tailwindcss/-/tailwindcss-4.1.14.tgz", diff --git a/package.json b/package.json index 676ef1cb77..911beb70c1 100644 --- a/package.json +++ b/package.json @@ -2125,6 +2125,7 @@ "ansi-to-html": "^0.6.7", "bootstrap": "^5.0.0", "bootstrap-less": "^3.3.8", + "clsx": "^2.1.1", "cross-fetch": "^3.1.5", "encoding": "^0.1.13", "fast-deep-equal": "^2.0.1", @@ -2163,6 +2164,7 @@ "strip-comments": "^2.0.1", "styled-components": "^5.2.1", "svg-to-pdfkit": "^0.1.8", + "tailwind-merge": "^3.3.1", "tcp-port-used": "^1.0.1", "tmp": "^0.2.4", "url-parse": "^1.5.10", diff --git a/src/webviews/extension-side/dataframe/dataframeController.ts b/src/webviews/extension-side/dataframe/dataframeController.ts index 122c0215da..2bb06f8a7c 100644 --- a/src/webviews/extension-side/dataframe/dataframeController.ts +++ b/src/webviews/extension-side/dataframe/dataframeController.ts @@ -1,11 +1,13 @@ import { injectable } from 'inversify'; import { commands, + env, l10n, NotebookEdit, type NotebookEditor, type NotebookRendererMessaging, notebooks, + Uri, window, workspace, WorkspaceEdit @@ -15,6 +17,7 @@ import type { IExtensionSyncActivationService } from '../../../platform/activati import type { IDisposable } from '../../../platform/common/types'; import { dispose } from '../../../platform/common/utils/lifecycle'; import { logger } from '../../../platform/logging'; +import { Output } from '../../../api'; type SelectPageSizeCommand = { cellId?: string; @@ -28,18 +31,29 @@ type GoToPageCommand = { page: number; }; -type CopyTableDataCommand = { +type CopyTableCommand = { cellId?: string; - command: 'copyTableData'; - data: string; + command: 'copyTable'; }; -type ExportDataframeCommand = { +type ExportTableCommand = { cellId?: string; - command: 'exportDataframe'; + command: 'exportTable'; }; -type DataframeCommand = SelectPageSizeCommand | GoToPageCommand | CopyTableDataCommand | ExportDataframeCommand; +interface DataFrameObject { + column_count: number; + columns: { + dtype: string; + name: string; + }[]; + preview_row_count: number; + row_count: number; + rows: Record[]; + type: string; +} + +type DataframeCommand = SelectPageSizeCommand | GoToPageCommand | CopyTableCommand | ExportTableCommand; @injectable() export class DataframeController implements IExtensionSyncActivationService { @@ -55,75 +69,177 @@ export class DataframeController implements IExtensionSyncActivationService { dispose(this.disposables); } - private async onDidReceiveMessage( - _comms: NotebookRendererMessaging, - { editor, message }: { editor: NotebookEditor; message: DataframeCommand } - ) { - logger.info('DataframeController received message', message); + private dataframeToCsv(dataframe: DataFrameObject): string { + console.log('Converting dataframe to CSV:', dataframe); - if (!message || typeof message !== 'object') { - return; + if (!dataframe || !dataframe.columns || !dataframe.rows) { + return ''; } - if (message.command === 'selectPageSize') { - return this.handleSelectPageSize(editor, message); - } + const headers = dataframe.columns; + const rows = dataframe.rows; + const columnNames = headers + .map((header: { name: string }) => header.name) + .filter((name: string | undefined) => Boolean(name)) + .filter((name: string) => !name.trim().toLowerCase().startsWith('_deepnote')) as string[]; - if (message.command === 'goToPage') { - return this.handleGoToPage(editor, message); - } + const csvRows = rows.map((row: Record) => columnNames.map((col) => row[col]).join(',')); - if (message.command === 'copyTableData') { - // TODO: Implement dataframe export functionality + return [columnNames.join(','), ...csvRows].join('\n'); + } + + private async getDataframeFromDataframeOutput(outputs: readonly Output[]): Promise { + if (outputs.length === 0) { + await this.showErrorToUser(l10n.t('No outputs found in the cell.')); return; } - if (message.command === 'exportDataframe') { - // TODO: Implement dataframe export functionality + const [firstOutput] = outputs; + + const item = firstOutput.items.find( + (i: { data: unknown; mime: string }) => i.mime === 'application/vnd.deepnote.dataframe.v3+json' + ); + + if (!item) { + await this.showErrorToUser( + l10n.t('No dataframe output found in the cell. Please ensure the cell has been executed.') + ); return; } - logger.warn(`DataframeController received unknown command:`, message); + const buffer = item.data as Uint8Array; + const json = new TextDecoder('utf-8').decode(buffer); + const dataframe = JSON.parse(json) as DataFrameObject; + + return dataframe; } - private async handleSelectPageSize(editor: NotebookEditor, message: SelectPageSizeCommand) { + private async handleCopyTable(editor: NotebookEditor, message: CopyTableCommand) { if (!message.cellId) { - const errorMessage = l10n.t( - 'Unable to update page size: No cell identifier provided. Please re-run the cell to update the output metadata.' + return this.showErrorToUser( + l10n.t( + 'Unable to copy table: No cell identifier provided. Please re-run the cell to update the output metadata.' + ) ); + } + + const cells = editor.notebook.getCells(); + const cell = cells.find((c) => c.metadata.id === message.cellId); + + if (!cell) { + return this.showErrorToUser( + l10n.t( + 'Unable to copy table: Could not find the cell with ID {0}. The cell may have been deleted.', + message.cellId ?? '' + ) + ); + } - logger.error(`[DataframeController] ${errorMessage}`); + const dataframe = await this.getDataframeFromDataframeOutput(cell.outputs); - await window.showErrorMessage(errorMessage); + if (!dataframe) { + return; + } + + const csv = this.dataframeToCsv(dataframe); + + if (!csv) { + return this.showErrorToUser(l10n.t('The dataframe is empty or could not be converted to CSV format.')); + } + + await env.clipboard.writeText(csv); + + await window.showInformationMessage(l10n.t('Text copied to clipboard!')); + + logger.info('[DataframeController] Dataframe copied to clipboard as CSV.'); + } - throw new Error(errorMessage); + private async handleExportTable(editor: NotebookEditor, message: ExportTableCommand) { + if (!message.cellId) { + return this.showErrorToUser( + l10n.t( + 'Unable to export table: No cell identifier provided. Please re-run the cell to update the output metadata.' + ) + ); } const cells = editor.notebook.getCells(); const cell = cells.find((c) => c.metadata.id === message.cellId); if (!cell) { - const errorMessage = l10n.t( - 'Unable to update page size: Could not find the cell with ID {0}. The cell may have been deleted.', - message.cellId ?? '' + return this.showErrorToUser( + l10n.t( + 'Unable to export table: Could not find the cell with ID {0}. The cell may have been deleted.', + message.cellId ?? '' + ) ); + } - logger.error(`[DataframeController] ${errorMessage}`); + const dataframe = await this.getDataframeFromDataframeOutput(cell.outputs); - await window.showErrorMessage(errorMessage); + if (!dataframe) { + return; + } - throw new Error(errorMessage); + const csv = this.dataframeToCsv(dataframe); + + if (!csv) { + return this.showErrorToUser(l10n.t('The dataframe is empty or could not be converted to CSV format.')); } - const cellIndex = cell.index; + const filename = `dataframe_${new Date().toISOString().replace(/[:.]/g, '-')}.csv`; - // Update page size in table state within cell metadata + try { + const uri = await window.showSaveDialog({ + defaultUri: Uri.file(filename), + filters: { + 'CSV files': ['csv'], + 'All files': ['*'] + } + }); + + if (uri) { + const encoder = new TextEncoder(); + + await workspace.fs.writeFile(uri, encoder.encode(csv)); + + await window.showInformationMessage(`File saved to ${uri}`); + } + } catch (error) { + await window.showErrorMessage(`Failed to save file: ${error}`); + } + } + + private async handleGoToPage(editor: NotebookEditor, message: GoToPageCommand) { + if (!message.cellId) { + return this.showErrorToUser( + l10n.t( + 'Unable to navigate to page: No cell identifier provided. Please re-run the cell to update the output metadata.' + ) + ); + } + + const cells = editor.notebook.getCells(); + const cell = cells.find((c) => c.metadata.id === message.cellId); + + if (!cell) { + return this.showErrorToUser( + l10n.t( + 'Unable to navigate to page: Could not find the cell with ID {0}. The cell may have been deleted.', + message.cellId ?? '' + ) + ); + } + + // Update page index in table state within cell metadata const existingTableState = cell.metadata.deepnote_table_state || {}; const updatedTableState = { ...existingTableState, - pageSize: message.size + pageIndex: message.page }; + const cellIndex = cell.index; + const edit = new WorkspaceEdit(); const notebookEdit = NotebookEdit.updateCellMetadata(cellIndex, { ...cell.metadata, @@ -134,8 +250,8 @@ export class DataframeController implements IExtensionSyncActivationService { await workspace.applyEdit(edit); - // Re-execute the cell to apply the new page size - logger.info(`[DataframeRenderer] Re-executing cell ${cellIndex} with new page size`); + // Re-execute the cell to apply the new page + logger.info(`[DataframeController] Re-executing cell ${cellIndex} with new page index ${message.page}`); await commands.executeCommand('notebook.cell.execute', { ranges: [{ start: cellIndex, end: cellIndex + 1 }], @@ -143,44 +259,36 @@ export class DataframeController implements IExtensionSyncActivationService { }); } - private async handleGoToPage(editor: NotebookEditor, message: GoToPageCommand) { + private async handleSelectPageSize(editor: NotebookEditor, message: SelectPageSizeCommand) { if (!message.cellId) { - const errorMessage = l10n.t( - 'Unable to navigate to page: No cell identifier provided. Please re-run the cell to update the output metadata.' + return this.showErrorToUser( + l10n.t( + 'Unable to update page size: No cell identifier provided. Please re-run the cell to update the output metadata.' + ) ); - - logger.error(`[DataframeController] ${errorMessage}`); - - await window.showErrorMessage(errorMessage); - - throw new Error(errorMessage); } const cells = editor.notebook.getCells(); const cell = cells.find((c) => c.metadata.id === message.cellId); if (!cell) { - const errorMessage = l10n.t( - 'Unable to navigate to page: Could not find the cell with ID {0}. The cell may have been deleted.', - message.cellId ?? '' + return this.showErrorToUser( + l10n.t( + 'Unable to update page size: Could not find the cell with ID {0}. The cell may have been deleted.', + message.cellId ?? '' + ) ); - - logger.error(`[DataframeController] ${errorMessage}`); - - await window.showErrorMessage(errorMessage); - - throw new Error(errorMessage); } - // Update page index in table state within cell metadata + const cellIndex = cell.index; + + // Update page size in table state within cell metadata const existingTableState = cell.metadata.deepnote_table_state || {}; const updatedTableState = { ...existingTableState, - pageIndex: message.page + pageSize: message.size }; - const cellIndex = cell.index; - const edit = new WorkspaceEdit(); const notebookEdit = NotebookEdit.updateCellMetadata(cellIndex, { ...cell.metadata, @@ -191,12 +299,49 @@ export class DataframeController implements IExtensionSyncActivationService { await workspace.applyEdit(edit); - // Re-execute the cell to apply the new page - logger.info(`[DataframeController] Re-executing cell ${cellIndex} with new page index ${message.page}`); + // Re-execute the cell to apply the new page size + logger.info(`[DataframeRenderer] Re-executing cell ${cellIndex} with new page size`); await commands.executeCommand('notebook.cell.execute', { ranges: [{ start: cellIndex, end: cellIndex + 1 }], document: editor.notebook.uri }); } + + private async onDidReceiveMessage( + _comms: NotebookRendererMessaging, + { editor, message }: { editor: NotebookEditor; message: DataframeCommand } + ) { + logger.info('DataframeController received message', message); + + if (!message || typeof message !== 'object') { + return; + } + + if (message.command === 'selectPageSize') { + return this.handleSelectPageSize(editor, message); + } + + if (message.command === 'goToPage') { + return this.handleGoToPage(editor, message); + } + + if (message.command === 'copyTable') { + return this.handleCopyTable(editor, message); + } + + if (message.command === 'exportTable') { + return this.handleExportTable(editor, message); + } + + logger.warn(`DataframeController received unknown command:`, message); + } + + private async showErrorToUser(errorMessage: string) { + logger.error(`[DataframeController] ${errorMessage}`); + + await window.showErrorMessage(errorMessage); + + throw new Error(errorMessage); + } } diff --git a/src/webviews/extension-side/dataframe/dataframeController.unit.test.ts b/src/webviews/extension-side/dataframe/dataframeController.unit.test.ts new file mode 100644 index 0000000000..7d30e60f90 --- /dev/null +++ b/src/webviews/extension-side/dataframe/dataframeController.unit.test.ts @@ -0,0 +1,742 @@ +import { assert } from 'chai'; +import { anything, instance, mock, when } from 'ts-mockito'; +import { + Disposable, + NotebookCell, + NotebookCellData, + NotebookCellKind, + NotebookCellOutput, + NotebookCellOutputItem, + NotebookDocument, + NotebookEditor, + NotebookRendererMessaging, + Uri +} from 'vscode'; +import { dispose } from '../../../platform/common/utils/lifecycle'; +import { IDisposable } from '../../../platform/common/types'; +import { DataframeController } from './dataframeController'; +import { createMockedNotebookDocument } from '../../../test/datascience/editor-integration/helpers'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; + +suite('DataframeController', () => { + let controller: DataframeController; + let disposables: IDisposable[] = []; + let comms: NotebookRendererMessaging; + let clipboard: any; + + setup(() => { + resetVSCodeMocks(); + disposables.push(new Disposable(() => resetVSCodeMocks())); + controller = new DataframeController(); + comms = mock(); + + // Get the mock clipboard instance from the env instance + clipboard = instance(mockedVSCodeNamespaces.env).clipboard; + + // Mock workspace.fs for file operations + const mockFs = { + // eslint-disable-next-line @typescript-eslint/no-empty-function + writeFile: async (_uri: Uri, _content: Uint8Array) => {} + }; + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(mockFs as any); + }); + + teardown(() => { + controller.dispose(); + disposables = dispose(disposables); + }); + + suite('CSV Conversion (dataframeToCsv)', () => { + test('Should return empty string for null dataframe', () => { + const result = (controller as any).dataframeToCsv(null); + assert.strictEqual(result, ''); + }); + + test('Should return empty string for undefined dataframe', () => { + const result = (controller as any).dataframeToCsv(undefined); + assert.strictEqual(result, ''); + }); + + test('Should return empty string for dataframe without columns', () => { + const dataframe = { + column_count: 0, + columns: undefined, + preview_row_count: 0, + row_count: 0, + rows: [], + type: 'dataframe' + }; + const result = (controller as any).dataframeToCsv(dataframe); + assert.strictEqual(result, ''); + }); + + test('Should return empty string for dataframe without rows', () => { + const dataframe = { + column_count: 2, + columns: [ + { dtype: 'int64', name: 'col1' }, + { dtype: 'string', name: 'col2' } + ], + preview_row_count: 0, + row_count: 0, + rows: undefined, + type: 'dataframe' + }; + const result = (controller as any).dataframeToCsv(dataframe); + assert.strictEqual(result, ''); + }); + + test('Should filter out columns starting with _deepnote', () => { + const dataframe = { + column_count: 3, + columns: [ + { dtype: 'int64', name: 'col1' }, + { dtype: 'string', name: '_deepnote_internal' }, + { dtype: 'string', name: 'col2' } + ], + preview_row_count: 2, + row_count: 2, + rows: [ + { col1: 1, _deepnote_internal: 'hidden', col2: 'a' }, + { col1: 2, _deepnote_internal: 'hidden', col2: 'b' } + ], + type: 'dataframe' + }; + const result = (controller as any).dataframeToCsv(dataframe); + const lines = result.split('\n'); + + assert.strictEqual(lines.length, 3); + assert.strictEqual(lines[0], 'col1,col2'); + assert.strictEqual(lines[1], '1,a'); + assert.strictEqual(lines[2], '2,b'); + }); + + test('Should properly format CSV with headers and data rows', () => { + const dataframe = { + column_count: 3, + columns: [ + { dtype: 'int64', name: 'id' }, + { dtype: 'string', name: 'name' }, + { dtype: 'float64', name: 'score' } + ], + preview_row_count: 2, + row_count: 2, + rows: [ + { id: 1, name: 'Alice', score: 95.5 }, + { id: 2, name: 'Bob', score: 87.3 } + ], + type: 'dataframe' + }; + const result = (controller as any).dataframeToCsv(dataframe); + const lines = result.split('\n'); + + assert.strictEqual(lines.length, 3); + assert.strictEqual(lines[0], 'id,name,score'); + assert.strictEqual(lines[1], '1,Alice,95.5'); + assert.strictEqual(lines[2], '2,Bob,87.3'); + }); + + test('Should handle empty column names', () => { + const dataframe = { + column_count: 3, + columns: [ + { dtype: 'int64', name: 'col1' }, + { dtype: 'string', name: '' }, + { dtype: 'string', name: 'col2' } + ], + preview_row_count: 1, + row_count: 1, + rows: [{ col1: 1, '': 'empty', col2: 'a' }], + type: 'dataframe' + }; + const result = (controller as any).dataframeToCsv(dataframe); + const lines = result.split('\n'); + + assert.strictEqual(lines[0], 'col1,col2'); + }); + + test('Should handle case-insensitive _deepnote filtering', () => { + const dataframe = { + column_count: 2, + columns: [ + { dtype: 'int64', name: 'col1' }, + { dtype: 'string', name: '_DEEPNOTE_test' } + ], + preview_row_count: 1, + row_count: 1, + rows: [{ col1: 1, _DEEPNOTE_test: 'hidden' }], + type: 'dataframe' + }; + const result = (controller as any).dataframeToCsv(dataframe); + const lines = result.split('\n'); + + assert.strictEqual(lines[0], 'col1'); + assert.strictEqual(lines[1], '1'); + }); + }); + + suite('Dataframe Extraction (getDataframeFromDataframeOutput)', () => { + test('Should show error for empty outputs array', async () => { + let errorShown = false; + when(mockedVSCodeNamespaces.window.showErrorMessage(anything())).thenCall(() => { + errorShown = true; + return Promise.resolve(); + }); + + try { + await (controller as any).getDataframeFromDataframeOutput([]); + assert.fail('Should have thrown an error'); + } catch (error) { + assert.isTrue(errorShown); + assert.include((error as Error).message, 'No outputs found'); + } + }); + + test('Should show error when dataframe MIME type not found', async () => { + let errorShown = false; + when(mockedVSCodeNamespaces.window.showErrorMessage(anything())).thenCall(() => { + errorShown = true; + return Promise.resolve(); + }); + + const outputs = [new NotebookCellOutput([NotebookCellOutputItem.text('some text', 'text/plain')])]; + + try { + await (controller as any).getDataframeFromDataframeOutput(outputs); + assert.fail('Should have thrown an error'); + } catch (error) { + assert.isTrue(errorShown); + assert.include((error as Error).message, 'No dataframe output found'); + } + }); + + test('Should successfully parse valid dataframe JSON', async () => { + const dataframeData = { + column_count: 2, + columns: [ + { dtype: 'int64', name: 'col1' }, + { dtype: 'string', name: 'col2' } + ], + preview_row_count: 1, + row_count: 1, + rows: [{ col1: 1, col2: 'test' }], + type: 'dataframe' + }; + + const jsonString = JSON.stringify(dataframeData); + const encoder = new TextEncoder(); + const buffer = encoder.encode(jsonString); + + const outputs = [ + new NotebookCellOutput([ + new NotebookCellOutputItem(buffer, 'application/vnd.deepnote.dataframe.v3+json') + ]) + ]; + + const result = await (controller as any).getDataframeFromDataframeOutput(outputs); + + assert.isDefined(result); + assert.deepStrictEqual(result, dataframeData); + }); + }); + + suite('Copy Table (handleCopyTable)', () => { + test('Should show error when cellId is missing', async () => { + let errorShown = false; + when(mockedVSCodeNamespaces.window.showErrorMessage(anything())).thenCall(() => { + errorShown = true; + throw new Error('No cell identifier'); + }); + + const editor = createMockEditor([]); + const message = { command: 'copyTable' as const }; + + try { + await (controller as any).handleCopyTable(editor, message); + } catch (e) { + // Expected + } + + assert.isTrue(errorShown); + }); + + test('Should show error when cell not found', async () => { + let errorShown = false; + when(mockedVSCodeNamespaces.window.showErrorMessage(anything())).thenCall(() => { + errorShown = true; + throw new Error('Cell not found'); + }); + + const cell = createCellWithOutputs('print(1)', [], { id: 'cell1' }); + const { editor } = createNotebookWithCell(cell); + + const message = { command: 'copyTable' as const, cellId: 'nonexistent' }; + + try { + await (controller as any).handleCopyTable(editor, message); + } catch (e) { + // Expected + } + + assert.isTrue(errorShown); + }); + + test('Should successfully copy dataframe to clipboard', async () => { + const dataframeData = { + column_count: 2, + columns: [ + { dtype: 'int64', name: 'id' }, + { dtype: 'string', name: 'name' } + ], + preview_row_count: 2, + row_count: 2, + rows: [ + { id: 1, name: 'Alice' }, + { id: 2, name: 'Bob' } + ], + type: 'dataframe' + }; + + const jsonString = JSON.stringify(dataframeData); + const encoder = new TextEncoder(); + const buffer = encoder.encode(jsonString); + + const cell = createCellWithOutputs( + 'df', + [ + new NotebookCellOutput([ + new NotebookCellOutputItem(buffer, 'application/vnd.deepnote.dataframe.v3+json') + ]) + ], + { id: 'cell1' } + ); + + const { editor } = createNotebookWithCell(cell); + + let messageShown = false; + + when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenCall(() => { + messageShown = true; + return Promise.resolve(); + }); + + const message = { command: 'copyTable' as const, cellId: 'cell1' }; + + await (controller as any).handleCopyTable(editor, message); + + const clipboardContent = await clipboard.readText(); + assert.strictEqual(clipboardContent, 'id,name\n1,Alice\n2,Bob'); + assert.isTrue(messageShown); + }); + + test('Should show error when dataframe is empty', async () => { + const dataframeData = { + column_count: 0, + columns: [], + preview_row_count: 0, + row_count: 0, + rows: [], + type: 'dataframe' + }; + + const jsonString = JSON.stringify(dataframeData); + const encoder = new TextEncoder(); + const buffer = encoder.encode(jsonString); + + const cell = createCellWithOutputs( + 'df', + [ + new NotebookCellOutput([ + new NotebookCellOutputItem(buffer, 'application/vnd.deepnote.dataframe.v3+json') + ]) + ], + { id: 'cell1' } + ); + + const { editor } = createNotebookWithCell(cell); + + let errorShown = false; + when(mockedVSCodeNamespaces.window.showErrorMessage(anything())).thenCall(() => { + errorShown = true; + throw new Error('Empty dataframe'); + }); + + const message = { command: 'copyTable' as const, cellId: 'cell1' }; + + try { + await (controller as any).handleCopyTable(editor, message); + } catch (e) { + // Expected + } + + assert.isTrue(errorShown); + }); + }); + + suite('Export Table (handleExportTable)', () => { + test('Should show error when cellId is missing', async () => { + let errorShown = false; + when(mockedVSCodeNamespaces.window.showErrorMessage(anything())).thenCall(() => { + errorShown = true; + throw new Error('No cell identifier'); + }); + + const editor = createMockEditor([]); + const message = { command: 'exportTable' as const }; + + try { + await (controller as any).handleExportTable(editor, message); + } catch (e) { + // Expected + } + + assert.isTrue(errorShown); + }); + + test('Should show error when cell not found', async () => { + let errorShown = false; + when(mockedVSCodeNamespaces.window.showErrorMessage(anything())).thenCall(() => { + errorShown = true; + throw new Error('Cell not found'); + }); + + const cell = createCellWithOutputs('print(1)', [], { id: 'cell1' }); + const { editor } = createNotebookWithCell(cell); + + const message = { command: 'exportTable' as const, cellId: 'nonexistent' }; + + try { + await (controller as any).handleExportTable(editor, message); + } catch (e) { + // Expected + } + + assert.isTrue(errorShown); + }); + + test('Should handle user canceling save dialog', async () => { + const dataframeData = { + column_count: 2, + columns: [ + { dtype: 'int64', name: 'id' }, + { dtype: 'string', name: 'name' } + ], + preview_row_count: 1, + row_count: 1, + rows: [{ id: 1, name: 'Alice' }], + type: 'dataframe' + }; + + const jsonString = JSON.stringify(dataframeData); + const encoder = new TextEncoder(); + const buffer = encoder.encode(jsonString); + + const cell = createCellWithOutputs( + 'df', + [ + new NotebookCellOutput([ + new NotebookCellOutputItem(buffer, 'application/vnd.deepnote.dataframe.v3+json') + ]) + ], + { id: 'cell1' } + ); + + const { editor } = createNotebookWithCell(cell); + + when(mockedVSCodeNamespaces.window.showSaveDialog(anything())).thenResolve(undefined); + + const message = { command: 'exportTable' as const, cellId: 'cell1' }; + + await (controller as any).handleExportTable(editor, message); + + // Should not throw, just return + }); + + test('Should successfully write file when user confirms', async () => { + const dataframeData = { + column_count: 2, + columns: [ + { dtype: 'int64', name: 'id' }, + { dtype: 'string', name: 'name' } + ], + preview_row_count: 2, + row_count: 2, + rows: [ + { id: 1, name: 'Alice' }, + { id: 2, name: 'Bob' } + ], + type: 'dataframe' + }; + + const jsonString = JSON.stringify(dataframeData); + const encoder = new TextEncoder(); + const buffer = encoder.encode(jsonString); + + const cell = createCellWithOutputs( + 'df', + [ + new NotebookCellOutput([ + new NotebookCellOutputItem(buffer, 'application/vnd.deepnote.dataframe.v3+json') + ]) + ], + { id: 'cell1' } + ); + + const { editor } = createNotebookWithCell(cell); + + const saveUri = Uri.file('/tmp/test.csv'); + let fileWritten = false; + let writtenContent = ''; + let messageShown = false; + + // Mock fs.writeFile to track calls + const mockFs = { + writeFile: async (_uri: Uri, content: Uint8Array) => { + fileWritten = true; + writtenContent = new TextDecoder().decode(content); + } + }; + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(mockFs as any); + + when(mockedVSCodeNamespaces.window.showSaveDialog(anything())).thenReturn(Promise.resolve(saveUri)); + when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenCall(() => { + messageShown = true; + return Promise.resolve(); + }); + + const message = { command: 'exportTable' as const, cellId: 'cell1' }; + + await (controller as any).handleExportTable(editor, message); + + assert.isTrue(fileWritten); + assert.strictEqual(writtenContent, 'id,name\n1,Alice\n2,Bob'); + assert.isTrue(messageShown); + }); + + test('Should show error message when file write fails', async () => { + const dataframeData = { + column_count: 2, + columns: [ + { dtype: 'int64', name: 'id' }, + { dtype: 'string', name: 'name' } + ], + preview_row_count: 1, + row_count: 1, + rows: [{ id: 1, name: 'Alice' }], + type: 'dataframe' + }; + + const jsonString = JSON.stringify(dataframeData); + const encoder = new TextEncoder(); + const buffer = encoder.encode(jsonString); + + const cell = createCellWithOutputs( + 'df', + [ + new NotebookCellOutput([ + new NotebookCellOutputItem(buffer, 'application/vnd.deepnote.dataframe.v3+json') + ]) + ], + { id: 'cell1' } + ); + + const { editor } = createNotebookWithCell(cell); + + const saveUri = Uri.file('/tmp/test.csv'); + let errorShown = false; + + // Mock fs.writeFile to throw an error + const mockFs = { + writeFile: async (_uri: Uri, _content: Uint8Array) => { + throw new Error('Permission denied'); + } + }; + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(mockFs as any); + + when(mockedVSCodeNamespaces.window.showSaveDialog(anything())).thenReturn(Promise.resolve(saveUri)); + when(mockedVSCodeNamespaces.window.showErrorMessage(anything())).thenCall(() => { + errorShown = true; + return Promise.resolve(); + }); + + const message = { command: 'exportTable' as const, cellId: 'cell1' }; + + await (controller as any).handleExportTable(editor, message); + + assert.isTrue(errorShown); + }); + + test('Should show error when dataframe is empty', async () => { + const dataframeData = { + column_count: 0, + columns: [], + preview_row_count: 0, + row_count: 0, + rows: [], + type: 'dataframe' + }; + + const jsonString = JSON.stringify(dataframeData); + const encoder = new TextEncoder(); + const buffer = encoder.encode(jsonString); + + const cell = createCellWithOutputs( + 'df', + [ + new NotebookCellOutput([ + new NotebookCellOutputItem(buffer, 'application/vnd.deepnote.dataframe.v3+json') + ]) + ], + { id: 'cell1' } + ); + + const { editor } = createNotebookWithCell(cell); + + let errorShown = false; + when(mockedVSCodeNamespaces.window.showErrorMessage(anything())).thenCall(() => { + errorShown = true; + throw new Error('Empty dataframe'); + }); + + const message = { command: 'exportTable' as const, cellId: 'cell1' }; + + try { + await (controller as any).handleExportTable(editor, message); + } catch (e) { + // Expected + } + + assert.isTrue(errorShown); + }); + }); + + suite('Message Routing (onDidReceiveMessage)', () => { + test('Should handle copyTable command', async () => { + const dataframeData = { + column_count: 1, + columns: [{ dtype: 'int64', name: 'col1' }], + preview_row_count: 1, + row_count: 1, + rows: [{ col1: 1 }], + type: 'dataframe' + }; + + const jsonString = JSON.stringify(dataframeData); + const encoder = new TextEncoder(); + const buffer = encoder.encode(jsonString); + + const cell = createCellWithOutputs( + 'df', + [ + new NotebookCellOutput([ + new NotebookCellOutputItem(buffer, 'application/vnd.deepnote.dataframe.v3+json') + ]) + ], + { id: 'cell1' } + ); + + const { editor } = createNotebookWithCell(cell); + + when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenResolve(); + + const message = { command: 'copyTable' as const, cellId: 'cell1' }; + + await (controller as any).onDidReceiveMessage(instance(comms), { editor, message }); + + const clipboardContent = await clipboard.readText(); + assert.strictEqual(clipboardContent, 'col1\n1'); + }); + + test('Should handle exportTable command', async () => { + const dataframeData = { + column_count: 1, + columns: [{ dtype: 'int64', name: 'col1' }], + preview_row_count: 1, + row_count: 1, + rows: [{ col1: 1 }], + type: 'dataframe' + }; + + const jsonString = JSON.stringify(dataframeData); + const encoder = new TextEncoder(); + const buffer = encoder.encode(jsonString); + + const cell = createCellWithOutputs( + 'df', + [ + new NotebookCellOutput([ + new NotebookCellOutputItem(buffer, 'application/vnd.deepnote.dataframe.v3+json') + ]) + ], + { id: 'cell1' } + ); + + const { editor } = createNotebookWithCell(cell); + + const saveUri = Uri.file('/tmp/test.csv'); + let exportHandled = false; + + // Mock fs.writeFile to track calls + const mockFs = { + writeFile: async (_uri: Uri, _content: Uint8Array) => { + exportHandled = true; + } + }; + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(mockFs as any); + + when(mockedVSCodeNamespaces.window.showSaveDialog(anything())).thenReturn(Promise.resolve(saveUri)); + when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenResolve(); + + const message = { command: 'exportTable' as const, cellId: 'cell1' }; + + await (controller as any).onDidReceiveMessage(instance(comms), { editor, message }); + + assert.isTrue(exportHandled); + }); + + test('Should ignore null or invalid messages', async () => { + const editor = createMockEditor([]); + + await (controller as any).onDidReceiveMessage(instance(comms), { editor, message: null }); + await (controller as any).onDidReceiveMessage(instance(comms), { editor, message: 'invalid' }); + + // Should not throw + }); + }); + + function createMockEditor(cells: NotebookCellData[]): NotebookEditor { + const notebook = createMockedNotebookDocument(cells); + return { + notebook + } as NotebookEditor; + } + + function createCellWithOutputs( + value: string, + outputs: NotebookCellOutput[], + metadata: Record = {} + ): NotebookCell { + const cell = mock(); + when(cell.outputs).thenReturn(outputs); + when(cell.metadata).thenReturn(metadata); + when(cell.kind).thenReturn(NotebookCellKind.Code); + when(cell.index).thenReturn(0); + const document = mock(); + when(document.getText()).thenReturn(value); + when(cell.document).thenReturn(instance(document)); + return instance(cell); + } + + function createNotebookWithCell(cell: NotebookCell): { notebook: NotebookDocument; editor: NotebookEditor } { + const notebook = mock(); + when(notebook.getCells()).thenReturn([cell]); + when(notebook.cellAt(0)).thenReturn(cell); + when(notebook.cellCount).thenReturn(1); + + const editor = { + notebook: instance(notebook) + } as NotebookEditor; + + return { notebook: instance(notebook), editor }; + } +}); diff --git a/src/webviews/webview-side/dataframe-renderer/DataframeRenderer.tsx b/src/webviews/webview-side/dataframe-renderer/DataframeRenderer.tsx index 8c8701beed..f2ad19ba45 100644 --- a/src/webviews/webview-side/dataframe-renderer/DataframeRenderer.tsx +++ b/src/webviews/webview-side/dataframe-renderer/DataframeRenderer.tsx @@ -1,8 +1,11 @@ -import React from 'react'; +import { clsx, type ClassValue } from 'clsx'; +import React, { ReactElement, ReactNode } from 'react'; import { memo, useMemo, useState } from 'react'; +import { twMerge } from 'tailwind-merge'; import type { RendererContext } from 'vscode-notebook-renderer'; import '../react-common/codicon/codicon.css'; +import { generateUuid } from '../../../platform/common/uuid'; export interface DataframeMetadata { table_state_spec?: string; @@ -68,9 +71,7 @@ export const DataframeRenderer = memo(function DataframeRenderer({ const [pageSize, setPageSize] = useState(tableState.pageSize || 10); const [pageIndex, setPageIndex] = useState(tableState.pageIndex || 0); - console.log( - `[DataframeRenderer] State: ${JSON.stringify(context.getState())}, tableState: ${JSON.stringify(tableState)}` - ); + const selectId = useMemo(() => generateUuid(), []); const filteredColumns = data.columns.filter((column) => !column.name.startsWith('_deepnote_')); const numberOfRows = Math.min(data.row_count, data.preview_row_count); @@ -112,6 +113,28 @@ export const DataframeRenderer = memo(function DataframeRenderer({ context.postMessage?.(message); }; + const handleCopyTable = () => { + console.log(`[DataframeRenderer] handleCopyTable called with cellId: ${cellId}`); + + const message = { + cellId, + command: 'copyTable' + }; + + context.postMessage?.(message); + }; + + const handleExportTable = () => { + console.log(`[DataframeRenderer] handleExportTable called with cellId: ${cellId}`); + + const message = { + cellId, + command: 'exportTable' + }; + + context.postMessage?.(message); + }; + return (
@@ -149,27 +172,20 @@ export const DataframeRenderer = memo(function DataframeRenderer({
{numberOfRows} rows, {numberOfColumns} columns
-
- - +
- + Page {pageIndex + 1} of {totalPages} - +
-
{/* Actions */}
+
+
+ +
+ + + +
+ +
+
); }); + +interface IconButtonProps extends React.ButtonHTMLAttributes { + children: ReactNode; +} + +function IconButton(props: IconButtonProps): ReactElement { + return ( + + ); +} + +function cn(...inputs: ClassValue[]) { + return twMerge(clsx(inputs)); +} From 4cc735b6f66dbedf226d6531ac72a4819d04302b Mon Sep 17 00:00:00 2001 From: Christoffer Artmann Date: Wed, 15 Oct 2025 14:58:12 +0200 Subject: [PATCH 2/2] pr feedback --- src/messageTypes.ts | 7 + src/platform/common/utils/localize.ts | 7 + src/platform/webviews/webviewHost.ts | 9 +- .../dataframe/dataframeController.ts | 58 ++++++-- .../dataframeController.unit.test.ts | 126 ++++++++++++++++++ .../dataframe-renderer/DataframeRenderer.tsx | 30 +++-- 6 files changed, 215 insertions(+), 22 deletions(-) diff --git a/src/messageTypes.ts b/src/messageTypes.ts index d456098903..273c8f7bb5 100644 --- a/src/messageTypes.ts +++ b/src/messageTypes.ts @@ -158,6 +158,13 @@ export type LocalizedMessages = { selectedImageListLabel: string; selectedImageLabel: string; dvDeprecationWarning: string; + dataframeRowsColumns: string; + dataframePerPage: string; + dataframePreviousPage: string; + dataframeNextPage: string; + dataframePageOf: string; + dataframeCopyTable: string; + dataframeExportTable: string; }; // Map all messages to specific payloads export class IInteractiveWindowMapping { diff --git a/src/platform/common/utils/localize.ts b/src/platform/common/utils/localize.ts index 75f75e7ac8..efab91f883 100644 --- a/src/platform/common/utils/localize.ts +++ b/src/platform/common/utils/localize.ts @@ -805,6 +805,13 @@ export namespace WebViews { export const dvDeprecationWarning = l10n.t( 'The built-in data viewer will be deprecated and no longer usable starting with Visual Studio Code 1.92. Please install other data viewing extensions to continue inspecting data' ); + export const dataframeRowsColumns = l10n.t('{0} rows, {1} columns'); + export const dataframePerPage = l10n.t('/ page'); + export const dataframePreviousPage = l10n.t('Previous page'); + export const dataframeNextPage = l10n.t('Next page'); + export const dataframePageOf = l10n.t('Page {0} of {1}'); + export const dataframeCopyTable = l10n.t('Copy table'); + export const dataframeExportTable = l10n.t('Export table'); } export namespace Deprecated { diff --git a/src/platform/webviews/webviewHost.ts b/src/platform/webviews/webviewHost.ts index b5ad7ce157..b26727cdfd 100644 --- a/src/platform/webviews/webviewHost.ts +++ b/src/platform/webviews/webviewHost.ts @@ -258,7 +258,14 @@ export abstract class WebviewHost implements IDisposable { deletePlot: localize.WebViews.deletePlot, selectedImageListLabel: localize.WebViews.selectedImageListLabel, selectedImageLabel: localize.WebViews.selectedImageLabel, - dvDeprecationWarning: localize.WebViews.dvDeprecationWarning + dvDeprecationWarning: localize.WebViews.dvDeprecationWarning, + dataframeRowsColumns: localize.WebViews.dataframeRowsColumns, + dataframePerPage: localize.WebViews.dataframePerPage, + dataframePreviousPage: localize.WebViews.dataframePreviousPage, + dataframeNextPage: localize.WebViews.dataframeNextPage, + dataframePageOf: localize.WebViews.dataframePageOf, + dataframeCopyTable: localize.WebViews.dataframeCopyTable, + dataframeExportTable: localize.WebViews.dataframeExportTable }; this.postMessageInternal(SharedMessages.LocInit, JSON.stringify(locStrings)).catch(noop); } diff --git a/src/webviews/extension-side/dataframe/dataframeController.ts b/src/webviews/extension-side/dataframe/dataframeController.ts index 2bb06f8a7c..cfbef322ca 100644 --- a/src/webviews/extension-side/dataframe/dataframeController.ts +++ b/src/webviews/extension-side/dataframe/dataframeController.ts @@ -4,6 +4,7 @@ import { env, l10n, NotebookEdit, + type NotebookCellOutput, type NotebookEditor, type NotebookRendererMessaging, notebooks, @@ -17,7 +18,6 @@ import type { IExtensionSyncActivationService } from '../../../platform/activati import type { IDisposable } from '../../../platform/common/types'; import { dispose } from '../../../platform/common/utils/lifecycle'; import { logger } from '../../../platform/logging'; -import { Output } from '../../../api'; type SelectPageSizeCommand = { cellId?: string; @@ -69,8 +69,35 @@ export class DataframeController implements IExtensionSyncActivationService { dispose(this.disposables); } + private escapeCsvField(value: unknown): string { + // Handle null/undefined as empty string + if (value === null || value === undefined) { + return ''; + } + + // Convert to string + const stringValue = String(value); + + // Check if field needs quoting (contains comma, quote, or newline) + const needsQuoting = + stringValue.includes(',') || + stringValue.includes('"') || + stringValue.includes('\n') || + stringValue.includes('\r'); + + if (needsQuoting) { + // Escape internal quotes by doubling them, then wrap in quotes + return `"${stringValue.replace(/"/g, '""')}"`; + } + + return stringValue; + } + private dataframeToCsv(dataframe: DataFrameObject): string { - console.log('Converting dataframe to CSV:', dataframe); + logger.debug('[DataframeController] Converting dataframe to CSV', { + columnCount: dataframe?.column_count, + rowCount: dataframe?.row_count + }); if (!dataframe || !dataframe.columns || !dataframe.rows) { return ''; @@ -83,20 +110,27 @@ export class DataframeController implements IExtensionSyncActivationService { .filter((name: string | undefined) => Boolean(name)) .filter((name: string) => !name.trim().toLowerCase().startsWith('_deepnote')) as string[]; - const csvRows = rows.map((row: Record) => columnNames.map((col) => row[col]).join(',')); + // Escape and join header names + const headerRow = columnNames.map((name) => this.escapeCsvField(name)).join(','); - return [columnNames.join(','), ...csvRows].join('\n'); + // Escape and join data rows + const csvRows = rows.map((row: Record) => + columnNames.map((col) => this.escapeCsvField(row[col])).join(',') + ); + + return [headerRow, ...csvRows].join('\n'); } - private async getDataframeFromDataframeOutput(outputs: readonly Output[]): Promise { + private async getDataframeFromDataframeOutput( + outputs: readonly NotebookCellOutput[] + ): Promise { if (outputs.length === 0) { await this.showErrorToUser(l10n.t('No outputs found in the cell.')); return; } - const [firstOutput] = outputs; - - const item = firstOutput.items.find( + const items = outputs.flatMap((output) => output.items); + const item = items.find( (i: { data: unknown; mime: string }) => i.mime === 'application/vnd.deepnote.dataframe.v3+json' ); @@ -203,10 +237,12 @@ export class DataframeController implements IExtensionSyncActivationService { await workspace.fs.writeFile(uri, encoder.encode(csv)); - await window.showInformationMessage(`File saved to ${uri}`); + await window.showInformationMessage(l10n.t('File saved to {0}', uri)); } } catch (error) { - await window.showErrorMessage(`Failed to save file: ${error}`); + const message = error instanceof Error ? error.message : String(error); + + await window.showErrorMessage(l10n.t('Failed to save file: {0}', message)); } } @@ -300,7 +336,7 @@ export class DataframeController implements IExtensionSyncActivationService { await workspace.applyEdit(edit); // Re-execute the cell to apply the new page size - logger.info(`[DataframeRenderer] Re-executing cell ${cellIndex} with new page size`); + logger.info(`[DataframeController] Re-executing cell ${cellIndex} with new page size`); await commands.executeCommand('notebook.cell.execute', { ranges: [{ start: cellIndex, end: cellIndex + 1 }], diff --git a/src/webviews/extension-side/dataframe/dataframeController.unit.test.ts b/src/webviews/extension-side/dataframe/dataframeController.unit.test.ts index 7d30e60f90..01ec3a10b6 100644 --- a/src/webviews/extension-side/dataframe/dataframeController.unit.test.ts +++ b/src/webviews/extension-side/dataframe/dataframeController.unit.test.ts @@ -173,6 +173,132 @@ suite('DataframeController', () => { assert.strictEqual(lines[0], 'col1'); assert.strictEqual(lines[1], '1'); }); + + test('Should escape values containing commas', () => { + const dataframe = { + column_count: 2, + columns: [ + { dtype: 'int64', name: 'id' }, + { dtype: 'string', name: 'address' } + ], + preview_row_count: 2, + row_count: 2, + rows: [ + { id: 1, address: '123 Main St, Apt 4' }, + { id: 2, address: 'New York, NY' } + ], + type: 'dataframe' + }; + const result = (controller as any).dataframeToCsv(dataframe); + const lines = result.split('\n'); + + assert.strictEqual(lines[0], 'id,address'); + assert.strictEqual(lines[1], '1,"123 Main St, Apt 4"'); + assert.strictEqual(lines[2], '2,"New York, NY"'); + }); + + test('Should escape values containing quotes', () => { + const dataframe = { + column_count: 2, + columns: [ + { dtype: 'int64', name: 'id' }, + { dtype: 'string', name: 'message' } + ], + preview_row_count: 2, + row_count: 2, + rows: [ + { id: 1, message: 'He said "Hello"' }, + { id: 2, message: 'She replied "Hi"' } + ], + type: 'dataframe' + }; + const result = (controller as any).dataframeToCsv(dataframe); + const lines = result.split('\n'); + + assert.strictEqual(lines[0], 'id,message'); + assert.strictEqual(lines[1], '1,"He said ""Hello"""'); + assert.strictEqual(lines[2], '2,"She replied ""Hi"""'); + }); + + test('Should escape values containing newlines', () => { + const dataframe = { + column_count: 2, + columns: [ + { dtype: 'int64', name: 'id' }, + { dtype: 'string', name: 'notes' } + ], + preview_row_count: 1, + row_count: 1, + rows: [{ id: 1, notes: 'Line 1\nLine 2' }], + type: 'dataframe' + }; + const result = (controller as any).dataframeToCsv(dataframe); + + // Don't split by \n since the quoted field contains newlines + assert.strictEqual(result, 'id,notes\n1,"Line 1\nLine 2"'); + }); + + test('Should handle null and undefined values', () => { + const dataframe = { + column_count: 3, + columns: [ + { dtype: 'int64', name: 'id' }, + { dtype: 'string', name: 'value1' }, + { dtype: 'string', name: 'value2' } + ], + preview_row_count: 2, + row_count: 2, + rows: [ + { id: 1, value1: null, value2: 'test' }, + { id: 2, value1: 'test', value2: undefined } + ], + type: 'dataframe' + }; + const result = (controller as any).dataframeToCsv(dataframe); + const lines = result.split('\n'); + + assert.strictEqual(lines[0], 'id,value1,value2'); + assert.strictEqual(lines[1], '1,,test'); + assert.strictEqual(lines[2], '2,test,'); + }); + + test('Should escape column names with special characters', () => { + const dataframe = { + column_count: 3, + columns: [ + { dtype: 'int64', name: 'id' }, + { dtype: 'string', name: 'name, title' }, + { dtype: 'string', name: 'description "quoted"' } + ], + preview_row_count: 1, + row_count: 1, + rows: [{ id: 1, 'name, title': 'John Doe', 'description "quoted"': 'Test' }], + type: 'dataframe' + }; + const result = (controller as any).dataframeToCsv(dataframe); + const lines = result.split('\n'); + + assert.strictEqual(lines[0], 'id,"name, title","description ""quoted"""'); + assert.strictEqual(lines[1], '1,John Doe,Test'); + }); + + test('Should handle complex escaping scenarios', () => { + const dataframe = { + column_count: 2, + columns: [ + { dtype: 'int64', name: 'id' }, + { dtype: 'string', name: 'data' } + ], + preview_row_count: 1, + row_count: 1, + rows: [{ id: 1, data: 'Value with "quotes", commas, and\nnewlines' }], + type: 'dataframe' + }; + const result = (controller as any).dataframeToCsv(dataframe); + + // Don't split by \n since the quoted field contains newlines + assert.strictEqual(result, 'id,data\n1,"Value with ""quotes"", commas, and\nnewlines"'); + }); }); suite('Dataframe Extraction (getDataframeFromDataframeOutput)', () => { diff --git a/src/webviews/webview-side/dataframe-renderer/DataframeRenderer.tsx b/src/webviews/webview-side/dataframe-renderer/DataframeRenderer.tsx index f2ad19ba45..4dd85fa068 100644 --- a/src/webviews/webview-side/dataframe-renderer/DataframeRenderer.tsx +++ b/src/webviews/webview-side/dataframe-renderer/DataframeRenderer.tsx @@ -6,6 +6,7 @@ import type { RendererContext } from 'vscode-notebook-renderer'; import '../react-common/codicon/codicon.css'; import { generateUuid } from '../../../platform/common/uuid'; +import { getLocString } from '../react-common/locReactSide'; export interface DataframeMetadata { table_state_spec?: string; @@ -170,7 +171,9 @@ export const DataframeRenderer = memo(function DataframeRenderer({
- {numberOfRows} rows, {numberOfColumns} columns + {getLocString('dataframeRowsColumns', `{0} rows, {1} columns`) + .replace('{0}', String(numberOfRows)) + .replace('{1}', String(numberOfColumns))}
- +
handlePageChange(pageIndex - 1)} >
- Page {pageIndex + 1} of {totalPages} + {getLocString('dataframePageOf', 'Page {0} of {1}') + .replace('{0}', String(pageIndex + 1)) + .replace('{1}', String(totalPages))} = totalPages - 1} - title="Next page" + title={getLocString('dataframeNextPage', 'Next page')} type="button" onClick={() => handlePageChange(pageIndex + 1)} > @@ -224,13 +229,18 @@ export const DataframeRenderer = memo(function DataframeRenderer({
- +