feat: Add commands to add Text blocks to notebooks (#286)

* feat: Add commands to add Text blocks to notebooks

* Minor changes

* Fix typo

* Minor changes

* Change picker items labels

Author: tkislan

package.json

@@ -230,6 +230,36 @@
                 "category": "Deepnote",
                 "icon": "$(add)"
             },
+            {
+                "command": "deepnote.addTextBlock",
+                "title": "%deepnote.commands.addTextBlock.title%",
+                "category": "Deepnote",
+                "icon": "$(text-size)"
+            },
+            {
+                "command": "deepnote.addTextBlockParagraph",
+                "title": "%deepnote.commands.addTextBlockParagraph.title%",
+                "category": "Deepnote",
+                "icon": "$(text-size)"
+            },
+            {
+                "command": "deepnote.addTextBlockHeading1",
+                "title": "%deepnote.commands.addTextBlockHeading1.title%",
+                "category": "Deepnote",
+                "icon": "$(text-size)"
+            },
+            {
+                "command": "deepnote.addTextBlockHeading2",
+                "title": "%deepnote.commands.addTextBlockHeading2.title%",
+                "category": "Deepnote",
+                "icon": "$(text-size)"
+            },
+            {
+                "command": "deepnote.addTextBlockHeading3",
+                "title": "%deepnote.commands.addTextBlockHeading3.title%",
+                "category": "Deepnote",
+                "icon": "$(text-size)"
+            },
             {
                 "command": "deepnote.newNotebook",
                 "title": "%deepnote.commands.newNotebook.title%",
@@ -1026,6 +1056,11 @@
                     "group": "navigation@12",
                     "when": "notebookType == 'deepnote'"
                 },
+                {
+                    "command": "deepnote.addTextBlock",
+                    "group": "navigation@13",
+                    "when": "notebookType == 'deepnote'"
+                },
                 {
                     "command": "deepnote.restartkernel",
                     "group": "navigation/execute@5",

package.nls.json

@@ -266,6 +266,11 @@
     "deepnote.commands.addInputDateRangeBlock.title": "Add Input Date Range Block",
     "deepnote.commands.addInputFileBlock.title": "Add Input File Block",
     "deepnote.commands.addButtonBlock.title": "Add Button Block",
+    "deepnote.commands.addTextBlock.title": "Add Text Block",
+    "deepnote.commands.addTextBlockParagraph.title": "Add Paragraph Block",
+    "deepnote.commands.addTextBlockHeading1.title": "Add Heading 1 Block",
+    "deepnote.commands.addTextBlockHeading2.title": "Add Heading 2 Block",
+    "deepnote.commands.addTextBlockHeading3.title": "Add Heading 3 Block",
     "deepnote.commands.newNotebook.title": "New Notebook",
     "deepnote.commands.renameProject.title": "Rename Project",
     "deepnote.commands.deleteProject.title": "Delete Project",

src/commands.ts

@@ -197,6 +197,11 @@ export interface ICommandNameArgumentTypeMapping {
     [DSCommands.AddInputDateRangeBlock]: [];
     [DSCommands.AddInputFileBlock]: [];
     [DSCommands.AddButtonBlock]: [];
+    [DSCommands.AddTextBlock]: [];
+    [DSCommands.AddTextBlockHeading1]: [];
+    [DSCommands.AddTextBlockHeading2]: [];
+    [DSCommands.AddTextBlockHeading3]: [];
+    [DSCommands.AddTextBlockParagraph]: [];
     [DSCommands.NewProject]: [];
     [DSCommands.ImportNotebook]: [];
     [DSCommands.ImportJupyterNotebook]: [];

src/notebooks/deepnote/deepnoteNotebookCommandListener.ts

@@ -8,7 +8,9 @@ import {
     NotebookRange,
     NotebookCell,
     NotebookEditorRevealType,
-    l10n
+    l10n,
+    QuickPickItem,
+    NotebookEditor
 } from 'vscode';
 import z from 'zod';
 
@@ -33,6 +35,7 @@ import {
     DeepnoteSqlMetadata
 } from './deepnoteSchemas';
 import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes';
+import { Pocket } from '../../platform/deepnote/pocket';
 
 export type InputBlockType =
     | 'input-text'
@@ -45,6 +48,10 @@ export type InputBlockType =
     | 'input-file'
     | 'button';
 
+export const TEXT_BLOCK_TYPES = ['text-cell-p', 'text-cell-h1', 'text-cell-h2', 'text-cell-h3'] as const;
+
+export type TextBlockType = (typeof TEXT_BLOCK_TYPES)[number];
+
 export function getInputBlockMetadata(blockType: InputBlockType, variableName: string) {
     const defaultInput = {
         deepnote_variable_name: variableName
@@ -181,6 +188,29 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation
         this.disposableRegistry.push(
             commands.registerCommand(Commands.AddButtonBlock, () => this.addInputBlock('button'))
         );
+        this.disposableRegistry.push(
+            commands.registerCommand(Commands.AddTextBlock, () => this.addTextBlockThroughPicker())
+        );
+        this.disposableRegistry.push(
+            commands.registerCommand(Commands.AddTextBlockHeading1, () =>
+                this.addTextBlockCommandHandler({ textBlockType: 'text-cell-h1' })
+            )
+        );
+        this.disposableRegistry.push(
+            commands.registerCommand(Commands.AddTextBlockHeading2, () =>
+                this.addTextBlockCommandHandler({ textBlockType: 'text-cell-h2' })
+            )
+        );
+        this.disposableRegistry.push(
+            commands.registerCommand(Commands.AddTextBlockHeading3, () =>
+                this.addTextBlockCommandHandler({ textBlockType: 'text-cell-h3' })
+            )
+        );
+        this.disposableRegistry.push(
+            commands.registerCommand(Commands.AddTextBlockParagraph, () =>
+                this.addTextBlockCommandHandler({ textBlockType: 'text-cell-p' })
+            )
+        );
     }
 
     public async addSqlBlock(): Promise<void> {
@@ -368,4 +398,92 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation
         // Enter edit mode on the new cell
         await commands.executeCommand('notebook.cell.edit');
     }
+
+    public async addTextBlockThroughPicker(): Promise<void> {
+        const TEXT_BLOCK_TYPE_LABELS = {
+            'text-cell-p': l10n.t('Paragraph'),
+            'text-cell-h1': l10n.t('Heading 1'),
+            'text-cell-h2': l10n.t('Heading 2'),
+            'text-cell-h3': l10n.t('Heading 3')
+        } as const satisfies Record<TextBlockType, string>;
+
+        const editor = window.activeNotebookEditor;
+        if (!editor) {
+            throw new Error(l10n.t('No active notebook editor found'));
+        }
+
+        const items: (QuickPickItem & { textBlockType: TextBlockType })[] = TEXT_BLOCK_TYPES.map((textBlockType) => {
+            const label = TEXT_BLOCK_TYPE_LABELS[textBlockType];
+            const description = l10n.t('Add a {0} text block', label);
+            return {
+                label,
+                description,
+                textBlockType
+            };
+        });
+
+        const selected = await window.showQuickPick(items, {
+            placeHolder: l10n.t('Select a text block type'),
+            matchOnDescription: true,
+            matchOnDetail: true
+        });
+
+        if (selected == null) {
+            return;
+        }
+
+        logger.info(`Selected text block type: ${selected.textBlockType}`);
+
+        await this.addTextBlock({ editor, textBlockType: selected.textBlockType });
+    }
+
+    public async addTextBlockCommandHandler({ textBlockType }: { textBlockType: TextBlockType }): Promise<void> {
+        const editor = window.activeNotebookEditor;
+        if (!editor) {
+            throw new Error(l10n.t('No active notebook editor found'));
+        }
+
+        await this.addTextBlock({ editor, textBlockType });
+    }
+
+    public async addTextBlock({
+        editor,
+        textBlockType
+    }: {
+        editor: NotebookEditor;
+        textBlockType: TextBlockType;
+    }): Promise<void> {
+        const TEXT_BLOCK_TYPE_EMPTY_VALUES = {
+            'text-cell-p': '',
+            'text-cell-h1': '# ',
+            'text-cell-h2': '## ',
+            'text-cell-h3': '### '
+        } as const satisfies Record<TextBlockType, string>;
+
+        const cellContent = TEXT_BLOCK_TYPE_EMPTY_VALUES[textBlockType];
+
+        const document = editor.notebook;
+        const selection = editor.selection;
+        const insertIndex = selection ? selection.end : document.cellCount;
+
+        const result = await notebookUpdaterUtils.chainWithPendingUpdates(document, (edit) => {
+            const newCell = new NotebookCellData(NotebookCellKind.Markup, cellContent, 'markdown');
+            newCell.metadata = {
+                __deepnotePocket: {
+                    type: textBlockType
+                } satisfies Pocket
+            };
+            const nbEdit = NotebookEdit.insertCells(insertIndex, [newCell]);
+            edit.set(document.uri, [nbEdit]);
+        });
+        if (result !== true) {
+            throw new Error(l10n.t('Failed to insert text block'));
+        }
+
+        const notebookRange = new NotebookRange(insertIndex, insertIndex + 1);
+        editor.revealRange(notebookRange, NotebookEditorRevealType.Default);
+        editor.selection = notebookRange;
+        // Enter edit mode on the new cell
+        await commands.executeCommand('notebook.cell.edit');
+    }
 }

src/platform/common/constants.ts

@@ -236,6 +236,11 @@ export namespace Commands {
     export const AddInputDateRangeBlock = 'deepnote.addInputDateRangeBlock';
     export const AddInputFileBlock = 'deepnote.addInputFileBlock';
     export const AddButtonBlock = 'deepnote.addButtonBlock';
+    export const AddTextBlock = 'deepnote.addTextBlock';
+    export const AddTextBlockParagraph = 'deepnote.addTextBlockParagraph';
+    export const AddTextBlockHeading1 = 'deepnote.addTextBlockHeading1';
+    export const AddTextBlockHeading2 = 'deepnote.addTextBlockHeading2';
+    export const AddTextBlockHeading3 = 'deepnote.addTextBlockHeading3';
     export const NewNotebook = 'deepnote.newNotebook';
     export const NewProject = 'deepnote.newProject';
     export const ImportNotebook = 'deepnote.importNotebook';