Skip to content

Commit a7cb1ca

Browse files
committed
feat: Add a command for adding chart blocks.
1 parent 065b4c0 commit a7cb1ca

File tree

6 files changed

+301
-4
lines changed

6 files changed

+301
-4
lines changed

package.json

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,12 @@
129129
"category": "Deepnote",
130130
"icon": "$(graph)"
131131
},
132+
{
133+
"command": "deepnote.addChartBlock",
134+
"title": "%deepnote.commands.addChartBlock.title%",
135+
"category": "Deepnote",
136+
"icon": "$(graph-line)"
137+
},
132138
{
133139
"command": "deepnote.addInputTextBlock",
134140
"title": "%deepnote.commands.addInputTextBlock.title%",
@@ -850,6 +856,51 @@
850856
"group": "navigation@1",
851857
"when": "notebookType == 'deepnote'"
852858
},
859+
{
860+
"command": "deepnote.addChartBlock",
861+
"group": "navigation@2",
862+
"when": "notebookType == 'deepnote'"
863+
},
864+
{
865+
"command": "deepnote.addBigNumberChartBlock",
866+
"group": "navigation@3",
867+
"when": "notebookType == 'deepnote'"
868+
},
869+
{
870+
"command": "deepnote.addInputTextBlock",
871+
"group": "navigation@4",
872+
"when": "notebookType == 'deepnote'"
873+
},
874+
{
875+
"command": "deepnote.addInputTextareaBlock",
876+
"group": "navigation@5",
877+
"when": "notebookType == 'deepnote'"
878+
},
879+
{
880+
"command": "deepnote.addInputSelectBlock",
881+
"group": "navigation@6",
882+
"when": "notebookType == 'deepnote'"
883+
},
884+
{
885+
"command": "deepnote.addInputSliderBlock",
886+
"group": "navigation@7",
887+
"when": "notebookType == 'deepnote'"
888+
},
889+
{
890+
"command": "deepnote.addInputCheckboxBlock",
891+
"group": "navigation@8",
892+
"when": "notebookType == 'deepnote'"
893+
},
894+
{
895+
"command": "deepnote.addInputDateBlock",
896+
"group": "navigation@9",
897+
"when": "notebookType == 'deepnote'"
898+
},
899+
{
900+
"command": "deepnote.addInputDateRangeBlock",
901+
"group": "navigation@10",
902+
"when": "notebookType == 'deepnote'"
903+
},
853904
{
854905
"command": "jupyter.restartkernel",
855906
"group": "navigation/execute@5",

package.nls.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -255,7 +255,8 @@
255255
"deepnote.commands.importNotebook.title": "Import Notebook",
256256
"deepnote.commands.importJupyterNotebook.title": "Import Jupyter Notebook",
257257
"deepnote.commands.addSqlBlock.title": "Add SQL Block",
258-
"deepnote.commands.addBigNumberChartBlock.title": "Add Big Number Chart Block",
258+
"deepnote.commands.addBigNumberChartBlock.title": "Add Big Number Block",
259+
"deepnote.commands.addChartBlock.title": "Add Chart Block",
259260
"deepnote.commands.addInputTextBlock.title": "Add Input Text Block",
260261
"deepnote.commands.addInputTextareaBlock.title": "Add Input Textarea Block",
261262
"deepnote.commands.addInputSelectBlock.title": "Add Input Select Block",

src/commands.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,7 @@ export interface ICommandNameArgumentTypeMapping {
188188
[DSCommands.ContinueEditSessionInCodespace]: [];
189189
[DSCommands.AddSqlBlock]: [];
190190
[DSCommands.AddBigNumberChartBlock]: [];
191+
[DSCommands.AddChartBlock]: [];
191192
[DSCommands.AddInputTextBlock]: [];
192193
[DSCommands.AddInputTextareaBlock]: [];
193194
[DSCommands.AddInputSelectBlock]: [];

src/notebooks/deepnote/deepnoteNotebookCommandListener.ts

Lines changed: 63 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,15 +75,18 @@ export function getInputBlockMetadata(blockType: InputBlockType, variableName: s
7575

7676
export function safeParseDeepnoteVariableNameFromContentJson(content: string): string | undefined {
7777
try {
78-
const variableNameResult = z.string().safeParse(JSON.parse(content)['deepnote_variable_name']);
78+
const parsed = JSON.parse(content);
79+
// Chart blocks use 'variable' key, other blocks use 'deepnote_variable_name'
80+
const variableName = parsed['variable'] ?? parsed['deepnote_variable_name'];
81+
const variableNameResult = z.string().safeParse(variableName);
7982
return variableNameResult.success ? variableNameResult.data : undefined;
8083
} catch (error) {
8184
logger.error('Error parsing deepnote variable name from content JSON', error);
8285
return undefined;
8386
}
8487
}
8588

86-
export function getNextDeepnoteVariableName(cells: NotebookCell[], prefix: 'df' | 'query' | 'input'): string {
89+
export function getNextDeepnoteVariableName(cells: NotebookCell[], prefix: 'df' | 'query' | 'input' | 'chart'): string {
8790
const deepnoteVariableNames = cells.reduce<string[]>((acc, cell) => {
8891
const contentValue = safeParseDeepnoteVariableNameFromContentJson(cell.document.getText());
8992

@@ -147,6 +150,7 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation
147150
this.disposableRegistry.push(
148151
commands.registerCommand(Commands.AddBigNumberChartBlock, () => this.addBigNumberChartBlock())
149152
);
153+
this.disposableRegistry.push(commands.registerCommand(Commands.AddChartBlock, () => this.addChartBlock()));
150154
this.disposableRegistry.push(
151155
commands.registerCommand(Commands.AddInputTextBlock, () => this.addInputBlock('input-text'))
152156
);
@@ -261,6 +265,63 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation
261265
await commands.executeCommand('notebook.cell.edit');
262266
}
263267

268+
public async addChartBlock(): Promise<void> {
269+
const editor = window.activeNotebookEditor;
270+
271+
if (!editor) {
272+
throw new Error(l10n.t('No active notebook editor found'));
273+
}
274+
275+
const document = editor.notebook;
276+
const selection = editor.selection;
277+
278+
const insertIndex = selection ? selection.end : document.cellCount;
279+
const deepnoteVariableName = getNextDeepnoteVariableName(document.getCells(), 'df');
280+
281+
const defaultVisualizationSpec = {
282+
mark: 'line',
283+
$schema: 'https://vega.github.io/schema/vega-lite/v5.json',
284+
data: { values: [] },
285+
encoding: {
286+
x: { field: 'x', type: 'quantitative' },
287+
y: { field: 'y', type: 'quantitative' }
288+
}
289+
};
290+
291+
const cellContent = {
292+
variable: deepnoteVariableName,
293+
spec: defaultVisualizationSpec,
294+
filters: []
295+
};
296+
297+
const metadata = {
298+
__deepnotePocket: {
299+
type: 'visualization'
300+
}
301+
};
302+
303+
const result = await chainWithPendingUpdates(document, (edit) => {
304+
const newCell = new NotebookCellData(NotebookCellKind.Code, JSON.stringify(cellContent, null, 2), 'json');
305+
306+
newCell.metadata = metadata;
307+
308+
const nbEdit = NotebookEdit.insertCells(insertIndex, [newCell]);
309+
310+
edit.set(document.uri, [nbEdit]);
311+
});
312+
313+
if (result !== true) {
314+
throw new Error(l10n.t('Failed to insert chart block'));
315+
}
316+
317+
const notebookRange = new NotebookRange(insertIndex, insertIndex + 1);
318+
319+
editor.revealRange(notebookRange, NotebookEditorRevealType.Default);
320+
editor.selection = notebookRange;
321+
322+
await commands.executeCommand('notebook.cell.edit');
323+
}
324+
264325
public async addInputBlock(blockType: InputBlockType): Promise<void> {
265326
const editor = window.activeNotebookEditor;
266327
if (!editor) {

src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts

Lines changed: 183 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ suite('DeepnoteNotebookCommandListener', () => {
102102
const TEST_INPUTS: Array<{
103103
description: string;
104104
cells: NotebookCell[];
105-
prefix: 'df' | 'query' | 'input';
105+
prefix: 'df' | 'query' | 'input' | 'chart';
106106
expected: string;
107107
}> = [
108108
// Tests with 'input' prefix
@@ -975,5 +975,187 @@ suite('DeepnoteNotebookCommandListener', () => {
975975
);
976976
});
977977
});
978+
979+
suite('addChartBlock', () => {
980+
test('should add chart block at the end when no selection exists', async () => {
981+
// Setup mocks
982+
const { editor, document } = createMockEditor([], undefined);
983+
const { chainStub, executeCommandStub, getCapturedNotebookEdits } =
984+
mockNotebookUpdateAndExecute(editor);
985+
986+
// Call the method
987+
await commandListener.addChartBlock();
988+
989+
const capturedNotebookEdits = getCapturedNotebookEdits();
990+
991+
// Verify chainWithPendingUpdates was called
992+
assert.isTrue(chainStub.calledOnce, 'chainWithPendingUpdates should be called once');
993+
assert.equal(chainStub.firstCall.args[0], document, 'Should be called with correct document');
994+
995+
// Verify the edits were captured
996+
assert.isNotNull(capturedNotebookEdits, 'Notebook edits should be captured');
997+
assert.isDefined(capturedNotebookEdits, 'Notebook edits should be defined');
998+
999+
const editsArray = capturedNotebookEdits!;
1000+
assert.equal(editsArray.length, 1, 'Should have one notebook edit');
1001+
1002+
const notebookEdit = editsArray[0] as any;
1003+
assert.equal(notebookEdit.newCells.length, 1, 'Should insert one cell');
1004+
1005+
const newCell = notebookEdit.newCells[0];
1006+
assert.equal(newCell.kind, NotebookCellKind.Code, 'Should be a code cell');
1007+
assert.equal(newCell.languageId, 'json', 'Should have json language');
1008+
1009+
// Verify cell content is valid JSON with correct structure
1010+
const content = JSON.parse(newCell.value);
1011+
assert.equal(content.variable, 'df_1', 'Should have correct variable name');
1012+
assert.property(content, 'spec', 'Should have spec property');
1013+
assert.property(content, 'filters', 'Should have filters property');
1014+
1015+
// Verify the spec has the correct Vega-Lite structure
1016+
assert.equal(content.spec.mark, 'line', 'Should be a line chart');
1017+
assert.equal(
1018+
content.spec.$schema,
1019+
'https://vega.github.io/schema/vega-lite/v5.json',
1020+
'Should have Vega-Lite schema'
1021+
);
1022+
assert.deepStrictEqual(content.spec.data, { values: [] }, 'Should have empty data array');
1023+
assert.property(content.spec, 'encoding', 'Should have encoding property');
1024+
assert.property(content.spec.encoding, 'x', 'Should have x encoding');
1025+
assert.property(content.spec.encoding, 'y', 'Should have y encoding');
1026+
1027+
// Verify metadata structure
1028+
assert.property(newCell.metadata, '__deepnotePocket', 'Should have __deepnotePocket metadata');
1029+
assert.equal(newCell.metadata.__deepnotePocket.type, 'visualization', 'Should have visualization type');
1030+
1031+
// Verify reveal and selection were set
1032+
assert.isTrue((editor.revealRange as sinon.SinonStub).calledOnce, 'Should reveal the new cell range');
1033+
const revealCall = (editor.revealRange as sinon.SinonStub).firstCall;
1034+
assert.equal(revealCall.args[0].start, 0, 'Should reveal correct range start');
1035+
assert.equal(revealCall.args[0].end, 1, 'Should reveal correct range end');
1036+
assert.equal(revealCall.args[1], 0, 'Should use NotebookEditorRevealType.Default (value 0)');
1037+
1038+
// Verify notebook.cell.edit command was executed
1039+
assert.isTrue(
1040+
executeCommandStub.calledWith('notebook.cell.edit'),
1041+
'Should execute notebook.cell.edit command'
1042+
);
1043+
});
1044+
1045+
test('should add chart block after selection when selection exists', async () => {
1046+
// Setup mocks
1047+
const existingCells = [createMockCell('{}'), createMockCell('{}')];
1048+
const selection = new NotebookRange(0, 1);
1049+
const { editor } = createMockEditor(existingCells, selection);
1050+
const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor);
1051+
1052+
// Call the method
1053+
await commandListener.addChartBlock();
1054+
1055+
const capturedNotebookEdits = getCapturedNotebookEdits();
1056+
1057+
// Verify chainWithPendingUpdates was called
1058+
assert.isTrue(chainStub.calledOnce, 'chainWithPendingUpdates should be called once');
1059+
1060+
// Verify a cell was inserted
1061+
assert.isNotNull(capturedNotebookEdits, 'Notebook edits should be captured');
1062+
const notebookEdit = capturedNotebookEdits![0] as any;
1063+
assert.equal(notebookEdit.newCells.length, 1, 'Should insert one cell');
1064+
assert.equal(notebookEdit.newCells[0].languageId, 'json', 'Should be JSON cell');
1065+
});
1066+
1067+
test('should generate correct variable name when existing chart variables exist', async () => {
1068+
// Setup mocks with existing df variables
1069+
const existingCells = [
1070+
createMockCell('{ "deepnote_variable_name": "df_1" }'),
1071+
createMockCell('{ "variable": "df_2" }')
1072+
];
1073+
const { editor } = createMockEditor(existingCells, undefined);
1074+
const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor);
1075+
1076+
// Call the method
1077+
await commandListener.addChartBlock();
1078+
1079+
const capturedNotebookEdits = getCapturedNotebookEdits();
1080+
const notebookEdit = capturedNotebookEdits![0] as any;
1081+
const newCell = notebookEdit.newCells[0];
1082+
1083+
// Verify variable name is df_3
1084+
const content = JSON.parse(newCell.value);
1085+
assert.equal(content.variable, 'df_3', 'Should generate next variable name');
1086+
});
1087+
1088+
test('should ignore other variable types when generating chart variable name', async () => {
1089+
// Setup mocks with input and df variables
1090+
const existingCells = [
1091+
createMockCell('{ "deepnote_variable_name": "input_10" }'),
1092+
createMockCell('{ "deepnote_variable_name": "df_5" }'),
1093+
createMockCell('{ "variable": "df_2" }')
1094+
];
1095+
const { editor } = createMockEditor(existingCells, undefined);
1096+
const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor);
1097+
1098+
// Call the method
1099+
await commandListener.addChartBlock();
1100+
1101+
const capturedNotebookEdits = getCapturedNotebookEdits();
1102+
const notebookEdit = capturedNotebookEdits![0] as any;
1103+
const newCell = notebookEdit.newCells[0];
1104+
1105+
// Verify variable name is df_6 (uses highest df_ number)
1106+
const content = JSON.parse(newCell.value);
1107+
assert.equal(content.variable, 'df_6', 'Should consider all df variables');
1108+
});
1109+
1110+
test('should insert at correct position in the middle of notebook', async () => {
1111+
// Setup mocks
1112+
const existingCells = [createMockCell('{}'), createMockCell('{}'), createMockCell('{}')];
1113+
const selection = new NotebookRange(1, 2);
1114+
const { editor } = createMockEditor(existingCells, selection);
1115+
const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor);
1116+
1117+
// Call the method
1118+
await commandListener.addChartBlock();
1119+
1120+
const capturedNotebookEdits = getCapturedNotebookEdits();
1121+
1122+
// Verify chainWithPendingUpdates was called
1123+
assert.isTrue(chainStub.calledOnce, 'chainWithPendingUpdates should be called once');
1124+
1125+
// Verify a cell was inserted
1126+
assert.isNotNull(capturedNotebookEdits, 'Notebook edits should be captured');
1127+
const notebookEdit = capturedNotebookEdits![0] as any;
1128+
assert.equal(notebookEdit.newCells.length, 1, 'Should insert one cell');
1129+
assert.equal(notebookEdit.newCells[0].languageId, 'json', 'Should be JSON cell');
1130+
});
1131+
1132+
test('should throw error when no active editor exists', async () => {
1133+
// Setup: no active editor
1134+
Object.defineProperty(window, 'activeNotebookEditor', {
1135+
value: undefined,
1136+
configurable: true,
1137+
writable: true
1138+
});
1139+
1140+
// Call the method and expect rejection
1141+
await assert.isRejected(commandListener.addChartBlock(), Error, 'No active notebook editor found');
1142+
});
1143+
1144+
test('should throw error when chainWithPendingUpdates fails', async () => {
1145+
// Setup mocks
1146+
const { editor } = createMockEditor([], undefined);
1147+
Object.defineProperty(window, 'activeNotebookEditor', {
1148+
value: editor,
1149+
configurable: true,
1150+
writable: true
1151+
});
1152+
1153+
// Mock chainWithPendingUpdates to return false
1154+
sandbox.stub(notebookUpdater, 'chainWithPendingUpdates').resolves(false);
1155+
1156+
// Call the method and expect rejection
1157+
await assert.isRejected(commandListener.addChartBlock(), Error, 'Failed to insert chart block');
1158+
});
1159+
});
9781160
});
9791161
});

src/platform/common/constants.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,7 @@ export namespace Commands {
227227
export const ManageIntegrations = 'deepnote.manageIntegrations';
228228
export const AddSqlBlock = 'deepnote.addSqlBlock';
229229
export const AddBigNumberChartBlock = 'deepnote.addBigNumberChartBlock';
230+
export const AddChartBlock = 'deepnote.addChartBlock';
230231
export const AddInputTextBlock = 'deepnote.addInputTextBlock';
231232
export const AddInputTextareaBlock = 'deepnote.addInputTextareaBlock';
232233
export const AddInputSelectBlock = 'deepnote.addInputSelectBlock';

0 commit comments

Comments
 (0)