diff --git a/package.json b/package.json index 3d3b59ec40..7b7ac852ef 100644 --- a/package.json +++ b/package.json @@ -129,6 +129,12 @@ "category": "Deepnote", "icon": "$(graph)" }, + { + "command": "deepnote.addChartBlock", + "title": "%deepnote.commands.addChartBlock.title%", + "category": "Deepnote", + "icon": "$(graph-line)" + }, { "command": "deepnote.addInputTextBlock", "title": "%deepnote.commands.addInputTextBlock.title%", @@ -850,6 +856,51 @@ "group": "navigation@1", "when": "notebookType == 'deepnote'" }, + { + "command": "deepnote.addChartBlock", + "group": "navigation@2", + "when": "notebookType == 'deepnote'" + }, + { + "command": "deepnote.addBigNumberChartBlock", + "group": "navigation@3", + "when": "notebookType == 'deepnote'" + }, + { + "command": "deepnote.addInputTextBlock", + "group": "navigation@4", + "when": "notebookType == 'deepnote'" + }, + { + "command": "deepnote.addInputTextareaBlock", + "group": "navigation@5", + "when": "notebookType == 'deepnote'" + }, + { + "command": "deepnote.addInputSelectBlock", + "group": "navigation@6", + "when": "notebookType == 'deepnote'" + }, + { + "command": "deepnote.addInputSliderBlock", + "group": "navigation@7", + "when": "notebookType == 'deepnote'" + }, + { + "command": "deepnote.addInputCheckboxBlock", + "group": "navigation@8", + "when": "notebookType == 'deepnote'" + }, + { + "command": "deepnote.addInputDateBlock", + "group": "navigation@9", + "when": "notebookType == 'deepnote'" + }, + { + "command": "deepnote.addInputDateRangeBlock", + "group": "navigation@10", + "when": "notebookType == 'deepnote'" + }, { "command": "jupyter.restartkernel", "group": "navigation/execute@5", diff --git a/package.nls.json b/package.nls.json index 1d6006666e..504e2565cf 100644 --- a/package.nls.json +++ b/package.nls.json @@ -255,7 +255,8 @@ "deepnote.commands.importNotebook.title": "Import Notebook", "deepnote.commands.importJupyterNotebook.title": "Import Jupyter Notebook", "deepnote.commands.addSqlBlock.title": "Add SQL Block", - "deepnote.commands.addBigNumberChartBlock.title": "Add Big Number Chart Block", + "deepnote.commands.addBigNumberChartBlock.title": "Add Big Number Block", + "deepnote.commands.addChartBlock.title": "Add Chart Block", "deepnote.commands.addInputTextBlock.title": "Add Input Text Block", "deepnote.commands.addInputTextareaBlock.title": "Add Input Textarea Block", "deepnote.commands.addInputSelectBlock.title": "Add Input Select Block", diff --git a/src/commands.ts b/src/commands.ts index 40192a8f6f..3cbc992794 100644 --- a/src/commands.ts +++ b/src/commands.ts @@ -188,6 +188,7 @@ export interface ICommandNameArgumentTypeMapping { [DSCommands.ContinueEditSessionInCodespace]: []; [DSCommands.AddSqlBlock]: []; [DSCommands.AddBigNumberChartBlock]: []; + [DSCommands.AddChartBlock]: []; [DSCommands.AddInputTextBlock]: []; [DSCommands.AddInputTextareaBlock]: []; [DSCommands.AddInputSelectBlock]: []; diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts index dcd086691f..a9bcc780a4 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts @@ -17,6 +17,7 @@ import { IExtensionSyncActivationService } from '../../platform/activation/types import { IDisposableRegistry } from '../../platform/common/types'; import { Commands } from '../../platform/common/constants'; import { chainWithPendingUpdates } from '../../kernels/execution/notebookUpdater'; +import { WrappedError } from '../../platform/errors/types'; import { DeepnoteBigNumberMetadataSchema, DeepnoteTextInputMetadataSchema, @@ -76,7 +77,10 @@ export function getInputBlockMetadata(blockType: InputBlockType, variableName: s export function safeParseDeepnoteVariableNameFromContentJson(content: string): string | undefined { try { - const variableNameResult = z.string().safeParse(JSON.parse(content)['deepnote_variable_name']); + const parsed = JSON.parse(content); + // Chart blocks use 'variable' key, other blocks use 'deepnote_variable_name' + const variableName = parsed['variable'] ?? parsed['deepnote_variable_name']; + const variableNameResult = z.string().safeParse(variableName); return variableNameResult.success ? variableNameResult.data : undefined; } catch (error) { logger.error('Error parsing deepnote variable name from content JSON', error); @@ -148,6 +152,7 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation this.disposableRegistry.push( commands.registerCommand(Commands.AddBigNumberChartBlock, () => this.addBigNumberChartBlock()) ); + this.disposableRegistry.push(commands.registerCommand(Commands.AddChartBlock, () => this.addChartBlock())); this.disposableRegistry.push( commands.registerCommand(Commands.AddInputTextBlock, () => this.addInputBlock('input-text')) ); @@ -262,6 +267,61 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation await commands.executeCommand('notebook.cell.edit'); } + public async addChartBlock(): Promise { + const editor = window.activeNotebookEditor; + + if (!editor) { + throw new WrappedError(l10n.t('No active notebook editor found')); + } + + const document = editor.notebook; + const selection = editor.selection; + const insertIndex = selection ? selection.end : document.cellCount; + + const defaultVisualizationSpec = { + mark: 'line', + $schema: 'https://vega.github.io/schema/vega-lite/v5.json', + data: { values: [] }, + encoding: { + x: { field: 'x', type: 'quantitative' }, + y: { field: 'y', type: 'quantitative' } + } + }; + + const cellContent = { + variable: 'df_1', + spec: defaultVisualizationSpec, + filters: [] + }; + + const metadata = { + __deepnotePocket: { + type: 'visualization' + } + }; + + const result = await chainWithPendingUpdates(document, (edit) => { + const newCell = new NotebookCellData(NotebookCellKind.Code, JSON.stringify(cellContent, null, 2), 'json'); + + newCell.metadata = metadata; + + const nbEdit = NotebookEdit.insertCells(insertIndex, [newCell]); + + edit.set(document.uri, [nbEdit]); + }); + + if (result !== true) { + throw new WrappedError(l10n.t('Failed to insert chart block')); + } + + const notebookRange = new NotebookRange(insertIndex, insertIndex + 1); + + editor.revealRange(notebookRange, NotebookEditorRevealType.Default); + editor.selection = notebookRange; + + await commands.executeCommand('notebook.cell.edit'); + } + public async addInputBlock(blockType: InputBlockType): Promise { const editor = window.activeNotebookEditor; if (!editor) { diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts index a282518292..05fd5e37cd 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts @@ -21,6 +21,7 @@ import { import { IDisposable } from '../../platform/common/types'; import * as notebookUpdater from '../../kernels/execution/notebookUpdater'; import { createMockedNotebookDocument } from '../../test/datascience/editor-integration/helpers'; +import { WrappedError } from '../../platform/errors/types'; import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes'; suite('DeepnoteNotebookCommandListener', () => { @@ -976,5 +977,191 @@ suite('DeepnoteNotebookCommandListener', () => { ); }); }); + + suite('addChartBlock', () => { + test('should add chart block at the end when no selection exists', async () => { + // Setup mocks + const { editor, document } = createMockEditor([], undefined); + const { chainStub, executeCommandStub, getCapturedNotebookEdits } = + mockNotebookUpdateAndExecute(editor); + + // Call the method + await commandListener.addChartBlock(); + + const capturedNotebookEdits = getCapturedNotebookEdits(); + + // Verify chainWithPendingUpdates was called + assert.isTrue(chainStub.calledOnce, 'chainWithPendingUpdates should be called once'); + assert.equal(chainStub.firstCall.args[0], document, 'Should be called with correct document'); + + // Verify the edits were captured + assert.isNotNull(capturedNotebookEdits, 'Notebook edits should be captured'); + assert.isDefined(capturedNotebookEdits, 'Notebook edits should be defined'); + + const editsArray = capturedNotebookEdits!; + assert.equal(editsArray.length, 1, 'Should have one notebook edit'); + + const notebookEdit = editsArray[0] as any; + assert.equal(notebookEdit.newCells.length, 1, 'Should insert one cell'); + + const newCell = notebookEdit.newCells[0]; + assert.equal(newCell.kind, NotebookCellKind.Code, 'Should be a code cell'); + assert.equal(newCell.languageId, 'json', 'Should have json language'); + + // Verify cell content is valid JSON with correct structure + const content = JSON.parse(newCell.value); + assert.equal(content.variable, 'df_1', 'Should have correct variable name'); + assert.property(content, 'spec', 'Should have spec property'); + assert.property(content, 'filters', 'Should have filters property'); + + // Verify the spec has the correct Vega-Lite structure + assert.equal(content.spec.mark, 'line', 'Should be a line chart'); + assert.equal( + content.spec.$schema, + 'https://vega.github.io/schema/vega-lite/v5.json', + 'Should have Vega-Lite schema' + ); + assert.deepStrictEqual(content.spec.data, { values: [] }, 'Should have empty data array'); + assert.property(content.spec, 'encoding', 'Should have encoding property'); + assert.property(content.spec.encoding, 'x', 'Should have x encoding'); + assert.property(content.spec.encoding, 'y', 'Should have y encoding'); + + // Verify metadata structure + assert.property(newCell.metadata, '__deepnotePocket', 'Should have __deepnotePocket metadata'); + assert.equal(newCell.metadata.__deepnotePocket.type, 'visualization', 'Should have visualization type'); + + // Verify reveal and selection were set + assert.isTrue((editor.revealRange as sinon.SinonStub).calledOnce, 'Should reveal the new cell range'); + const revealCall = (editor.revealRange as sinon.SinonStub).firstCall; + assert.equal(revealCall.args[0].start, 0, 'Should reveal correct range start'); + assert.equal(revealCall.args[0].end, 1, 'Should reveal correct range end'); + assert.equal(revealCall.args[1], 0, 'Should use NotebookEditorRevealType.Default (value 0)'); + + // Verify notebook.cell.edit command was executed + assert.isTrue( + executeCommandStub.calledWith('notebook.cell.edit'), + 'Should execute notebook.cell.edit command' + ); + }); + + test('should add chart block after selection when selection exists', async () => { + // Setup mocks + const existingCells = [createMockCell('{}'), createMockCell('{}')]; + const selection = new NotebookRange(0, 1); + const { editor } = createMockEditor(existingCells, selection); + const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); + + // Call the method + await commandListener.addChartBlock(); + + const capturedNotebookEdits = getCapturedNotebookEdits(); + + // Verify chainWithPendingUpdates was called + assert.isTrue(chainStub.calledOnce, 'chainWithPendingUpdates should be called once'); + + // Verify a cell was inserted + assert.isNotNull(capturedNotebookEdits, 'Notebook edits should be captured'); + const notebookEdit = capturedNotebookEdits![0] as any; + assert.equal(notebookEdit.newCells.length, 1, 'Should insert one cell'); + assert.equal(notebookEdit.newCells[0].languageId, 'json', 'Should be JSON cell'); + }); + + test('should use hardcoded variable name df_1', async () => { + // Setup mocks with existing df variables + const existingCells = [ + createMockCell('{ "deepnote_variable_name": "df_1" }'), + createMockCell('{ "variable": "df_2" }') + ]; + const { editor } = createMockEditor(existingCells, undefined); + const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); + + // Call the method + await commandListener.addChartBlock(); + + const capturedNotebookEdits = getCapturedNotebookEdits(); + const notebookEdit = capturedNotebookEdits![0] as any; + const newCell = notebookEdit.newCells[0]; + + // Verify variable name is always df_1 + const content = JSON.parse(newCell.value); + assert.equal(content.variable, 'df_1', 'Should always use df_1'); + }); + + test('should always use df_1 regardless of existing variables', async () => { + // Setup mocks with various existing variables + const existingCells = [ + createMockCell('{ "deepnote_variable_name": "input_10" }'), + createMockCell('{ "deepnote_variable_name": "df_5" }'), + createMockCell('{ "variable": "df_2" }') + ]; + const { editor } = createMockEditor(existingCells, undefined); + const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); + + // Call the method + await commandListener.addChartBlock(); + + const capturedNotebookEdits = getCapturedNotebookEdits(); + const notebookEdit = capturedNotebookEdits![0] as any; + const newCell = notebookEdit.newCells[0]; + + // Verify variable name is always df_1 + const content = JSON.parse(newCell.value); + assert.equal(content.variable, 'df_1', 'Should always use df_1'); + }); + + test('should insert at correct position in the middle of notebook', async () => { + // Setup mocks + const existingCells = [createMockCell('{}'), createMockCell('{}'), createMockCell('{}')]; + const selection = new NotebookRange(1, 2); + const { editor } = createMockEditor(existingCells, selection); + const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); + + // Call the method + await commandListener.addChartBlock(); + + const capturedNotebookEdits = getCapturedNotebookEdits(); + + // Verify chainWithPendingUpdates was called + assert.isTrue(chainStub.calledOnce, 'chainWithPendingUpdates should be called once'); + + // Verify a cell was inserted + assert.isNotNull(capturedNotebookEdits, 'Notebook edits should be captured'); + const notebookEdit = capturedNotebookEdits![0] as any; + assert.equal(notebookEdit.newCells.length, 1, 'Should insert one cell'); + assert.equal(notebookEdit.newCells[0].languageId, 'json', 'Should be JSON cell'); + }); + + test('should throw error when no active editor exists', async () => { + // Setup: no active editor + Object.defineProperty(window, 'activeNotebookEditor', { + value: undefined, + configurable: true, + writable: true + }); + + // Call the method and expect rejection + await assert.isRejected( + commandListener.addChartBlock(), + WrappedError, + 'No active notebook editor found' + ); + }); + + test('should throw error when chainWithPendingUpdates fails', async () => { + // Setup mocks + const { editor } = createMockEditor([], undefined); + Object.defineProperty(window, 'activeNotebookEditor', { + value: editor, + configurable: true, + writable: true + }); + + // Mock chainWithPendingUpdates to return false + sandbox.stub(notebookUpdater, 'chainWithPendingUpdates').resolves(false); + + // Call the method and expect rejection + await assert.isRejected(commandListener.addChartBlock(), WrappedError, 'Failed to insert chart block'); + }); + }); }); }); diff --git a/src/platform/common/constants.ts b/src/platform/common/constants.ts index c6010d716f..59e2b63e1f 100644 --- a/src/platform/common/constants.ts +++ b/src/platform/common/constants.ts @@ -227,6 +227,7 @@ export namespace Commands { export const ManageIntegrations = 'deepnote.manageIntegrations'; export const AddSqlBlock = 'deepnote.addSqlBlock'; export const AddBigNumberChartBlock = 'deepnote.addBigNumberChartBlock'; + export const AddChartBlock = 'deepnote.addChartBlock'; export const AddInputTextBlock = 'deepnote.addInputTextBlock'; export const AddInputTextareaBlock = 'deepnote.addInputTextareaBlock'; export const AddInputSelectBlock = 'deepnote.addInputSelectBlock';