add run above and run below to native notebooks (#4240)

* add run above and run below to native notebooks

* address some comments

* fix tests

* add tests

* register commands on the
nativeEditorCommandListener

* skipt tests

package.json

@@ -279,6 +279,26 @@
                 },
                 "enablement": "notebookViewType == jupyter-notebook && jupyter.isnotebooktrusted"
             },
+            {
+                "command": "jupyter.notebookeditor.runallcellsabove",
+                "title": "%DataScience.runAbove%",
+                "category": "Jupyter",
+                "icon": {
+                    "light": "resources/light/runabove.svg",
+                    "dark": "resources/dark/runabove-inverse.svg"
+                },
+                "enablement": "notebookViewType == jupyter-notebook && jupyter.isnotebooktrusted"
+            },
+            {
+                "command": "jupyter.notebookeditor.runcellandallbelow",
+                "title": "%DataScience.runBelow%",
+                "category": "Jupyter",
+                "icon": {
+                    "light": "resources/light/runbelow.svg",
+                    "dark": "resources/dark/runbelow-inverse.svg"
+                },
+                "enablement": "notebookViewType == jupyter-notebook && jupyter.isnotebooktrusted"
+            },
             {
                 "command": "jupyter.export",
                 "title": "%DataScience.notebookExportAs%",
@@ -733,6 +753,18 @@
                 }
             ],
             "editor/title": [
+                {
+                    "command": "jupyter.notebookeditor.runallcellsabove",
+                    "title": "%DataScience.runAbove%",
+                    "group": "navigation",
+                    "when": "resourceLangId == jupyter && notebookViewType == 'jupyter-notebook' && jupyter.isnotebooktrusted && jupyter.havenativecells"
+                },
+                {
+                    "command": "jupyter.notebookeditor.runcellandallbelow",
+                    "title": "%DataScience.runBelow%",
+                    "group": "navigation",
+                    "when": "resourceLangId == jupyter && notebookViewType == 'jupyter-notebook' && jupyter.isnotebooktrusted && jupyter.havenativecells"
+                },
                 {
                     "command": "jupyter.notebookeditor.restartkernel",
                     "title": "%jupyter.command.jupyter.restartkernel.title%",

package.nls.json

@@ -394,6 +394,8 @@
     "DataScience.insertAbove": "Insert cell above",
     "DataScience.addCell": "Add cell",
     "DataScience.runAll": "Run all cells",
+    "DataScience.runAbove": "Run cells above",
+    "DataScience.runBelow": "Run cell and below",
     "DataScience.convertingToPythonFile": "Converting ipynb to python file",
     "DataScience.untitledNotebookMessage": "Your changes will be lost if you don't save them.",
     "DataScience.untitledNotebookYes": "Save",

src/client/common/application/commands.ts

@@ -92,6 +92,8 @@ export interface ICommandNameArgumentTypeMapping extends ICommandNameWithoutArgu
     [DSCommands.RunCell]: [Uri, number, number, number, number];
     [DSCommands.RunAllCellsAbove]: [Uri, number, number];
     [DSCommands.RunCellAndAllBelow]: [Uri, number, number];
+    [DSCommands.NativeNotebookRunAllCellsAbove]: [Uri];
+    [DSCommands.NativeNotebookRunCellAndAllBelow]: [Uri];
     [DSCommands.RunAllCellsAbovePalette]: [];
     [DSCommands.RunCellAndAllBelowPalette]: [];
     [DSCommands.DebugCurrentCellPalette]: [];

src/client/datascience/constants.ts

@@ -49,6 +49,8 @@ export namespace Commands {
     export const RunAllCells = 'jupyter.runallcells';
     export const RunAllCellsAbove = 'jupyter.runallcellsabove';
     export const RunCellAndAllBelow = 'jupyter.runcellandallbelow';
+    export const NativeNotebookRunAllCellsAbove = 'jupyter.notebookeditor.runallcellsabove';
+    export const NativeNotebookRunCellAndAllBelow = 'jupyter.notebookeditor.runcellandallbelow';
     export const SetJupyterKernel = 'jupyter.setKernel';
     export const SwitchJupyterKernel = 'jupyter.switchKernel';
     export const RunAllCellsAbovePalette = 'jupyter.runallcellsabove.palette';

src/client/datascience/interactive-ipynb/nativeEditor.ts

@@ -418,6 +418,13 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
         throw Error('Not implemented Exception');
     }
 
+    public runAbove(): void {
+        throw Error('Not implemented Exception');
+    }
+    public runCellAndBelow(): void {
+        throw Error('Not implemented Exception');
+    }
+
     protected addSysInfo(reason: SysInfoReason): Promise<void> {
         // We need to send a message when restarting
         if (reason === SysInfoReason.Restart || reason === SysInfoReason.New) {

src/client/datascience/interactive-ipynb/nativeEditorCommandListener.ts

@@ -55,6 +55,14 @@ export class NativeEditorCommandListener implements IDataScienceCommandListener
         this.disposableRegistry.push(
             commandManager.registerCommand(Commands.NotebookEditorAddCellBelow, () => this.addCellBelow())
         );
+        this.disposableRegistry.push(
+            commandManager.registerCommand(Commands.NativeNotebookRunAllCellsAbove, (uri) => this.runAbove(uri))
+        );
+        this.disposableRegistry.push(
+            commandManager.registerCommand(Commands.NativeNotebookRunCellAndAllBelow, (uri) =>
+                this.runCellAndBelow(uri)
+            )
+        );
     }
 
     private runAllCells() {
@@ -130,4 +138,17 @@ export class NativeEditorCommandListener implements IDataScienceCommandListener
             }
         }
     }
+
+    private runAbove(uri: Uri): void {
+        const activeEditor = this.provider.activeEditor;
+        if (activeEditor) {
+            activeEditor.runAbove(uri);
+        }
+    }
+    private runCellAndBelow(uri: Uri): void {
+        const activeEditor = this.provider.activeEditor;
+        if (activeEditor) {
+            activeEditor.runCellAndBelow(uri);
+        }
+    }
 }

src/client/datascience/notebook/notebookEditor.ts

@@ -4,7 +4,7 @@
 'use strict';
 
 import { ConfigurationTarget, Event, EventEmitter, Uri, WebviewPanel } from 'vscode';
-import type { NotebookCell, NotebookDocument } from '../../../../types/vscode-proposed';
+import { NotebookCell, NotebookDocument } from '../../../../types/vscode-proposed';
 import { IApplicationShell, ICommandManager, IVSCodeNotebook } from '../../common/application/types';
 import { traceError } from '../../common/logger';
 import { IConfigurationService, IDisposable, IDisposableRegistry } from '../../common/types';
@@ -28,7 +28,6 @@ import { chainWithPendingUpdates } from './helpers/notebookUpdater';
 const vscodeNotebookEnums = require('vscode') as typeof import('vscode-proposed');
 
 export class NotebookEditor implements INotebookEditor {
-    public readonly type = 'native';
     public get onDidChangeViewState(): Event<void> {
         return this.changedViewState.event;
     }
@@ -63,6 +62,7 @@ export class NotebookEditor implements INotebookEditor {
     public get onExecutedCode(): Event<string> {
         return this.executedCode.event;
     }
+    public readonly type = 'native';
     public notebook?: INotebook | undefined;
 
     private changedViewState = new EventEmitter<void>();
@@ -256,6 +256,51 @@ export class NotebookEditor implements INotebookEditor {
         this._closed.fire(this);
     }
 
+    public runAbove(uri: Uri): void {
+        const cellId = this.getSelectedCellId(uri);
+        const index = this.document.cells.findIndex((c) => c.uri.toString() === cellId);
+
+        if (index > 0) {
+            // Get all cellIds until `index`.
+            const cells = this.document.cells.slice(0, index).map((cell) => cell);
+            this.runCellRange(cells);
+        }
+    }
+    public runCellAndBelow(uri: Uri): void {
+        const cellId = this.getSelectedCellId(uri);
+        const index = this.document.cells.findIndex((c) => c.uri.toString() === cellId);
+
+        if (index >= 0) {
+            // Get all cellIds starting from `index`.
+            const cells = this.document.cells.slice(index).map((cell) => cell);
+            this.runCellRange(cells);
+        }
+    }
+
+    private getSelectedCellId(uri: Uri): string | undefined {
+        const editor = this.vscodeNotebook.notebookEditors.find((nb) => nb.document.uri.toString() === uri.toString());
+
+        if (editor && editor.selection) {
+            return editor.selection.uri.toString();
+        }
+
+        return undefined;
+    }
+
+    private runCellRange(cells: NotebookCell[]) {
+        const kernel = this.kernelProvider.get(this.file);
+
+        if (!kernel || this.restartingKernel) {
+            return;
+        }
+
+        cells.forEach(async (cell) => {
+            if (cell.cellKind === vscodeNotebookEnums.CellKind.Code) {
+                await kernel.executeCell(cell);
+            }
+        });
+    }
+
     private async restartKernelInternal(kernel: IKernel): Promise<void> {
         this.restartingKernel = true;
 

src/client/datascience/types.ts

@@ -608,6 +608,8 @@ export interface INotebookEditor extends Disposable, IInteractiveBase {
     interruptKernel(): Promise<void>;
     restartKernel(): Promise<void>;
     syncAllCells(): Promise<void>;
+    runAbove(uri: Uri): void;
+    runCellAndBelow(uri: Uri): void;
 }
 
 export const INotebookExtensibility = Symbol('INotebookExtensibility');

src/test/datascience/notebook/notebookEditor.vscode.test.ts

@@ -0,0 +1,102 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+import { assert } from 'chai';
+import { ICommandManager, IVSCodeNotebook } from '../../../client/common/application/types';
+import { IDisposable } from '../../../client/common/types';
+import { Commands } from '../../../client/datascience/constants';
+import { INotebookEditorProvider } from '../../../client/datascience/types';
+import { IExtensionTestApi } from '../../common';
+import { initialize } from '../../initialize';
+import {
+    canRunNotebookTests,
+    closeNotebooks,
+    closeNotebooksAndCleanUpAfterTests,
+    deleteAllCellsAndWait,
+    insertCodeCell,
+    waitForExecutionCompletedSuccessfully,
+    waitForKernelToGetAutoSelected
+} from './helper';
+const vscodeNotebookEnums = require('vscode') as typeof import('vscode-proposed');
+
+suite('Notebook Editor tests', () => {
+    let api: IExtensionTestApi;
+    let vscodeNotebook: IVSCodeNotebook;
+    let editorProvider: INotebookEditorProvider;
+    let commandManager: ICommandManager;
+    const disposables: IDisposable[] = [];
+
+    suiteSetup(async function () {
+        api = await initialize();
+        if (!(await canRunNotebookTests())) {
+            return this.skip();
+        }
+        vscodeNotebook = api.serviceContainer.get<IVSCodeNotebook>(IVSCodeNotebook);
+        editorProvider = api.serviceContainer.get<INotebookEditorProvider>(INotebookEditorProvider);
+        commandManager = api.serviceContainer.get<ICommandManager>(ICommandManager);
+    });
+
+    setup(async function () {
+        // Open a notebook and use this for all tests in this test suite.
+        await editorProvider.createNew();
+        await waitForKernelToGetAutoSelected();
+        await deleteAllCellsAndWait();
+        assert.isOk(vscodeNotebook.activeNotebookEditor, 'No active notebook');
+    });
+
+    teardown(async function () {
+        await closeNotebooks(disposables);
+        await closeNotebooksAndCleanUpAfterTests(disposables);
+    });
+    suiteTeardown(() => closeNotebooksAndCleanUpAfterTests(disposables));
+
+    test('Run cells below', async function () {
+        // add some cells
+        // https://github.com/microsoft/vscode-jupyter/issues/4250
+        this.skip();
+        await insertCodeCell('print("0")', { index: 0 });
+        await insertCodeCell('print("1")', { index: 1 });
+        await insertCodeCell('print("2")', { index: 2 });
+
+        // run command
+        await commandManager.executeCommand(
+            Commands.NativeNotebookRunCellAndAllBelow,
+            vscodeNotebook.activeNotebookEditor?.document.uri!
+        );
+        const thirdCell = vscodeNotebook.activeNotebookEditor?.document.cells![2]!;
+        await waitForExecutionCompletedSuccessfully(thirdCell);
+
+        // The third cell should have a runState of Success
+        assert.strictEqual(thirdCell?.metadata.runState, vscodeNotebookEnums.NotebookCellRunState.Success);
+    });
+
+    test('Run cells above', async function () {
+        // This test is skipped because there is no way of selecting a cell in this context
+        // since by default the first cell is selected nothing happens when running all cells above
+        // https://github.com/microsoft/vscode-jupyter/issues/4250
+        this.skip();
+        // add some cells
+        await insertCodeCell('print("0")', { index: 0 });
+        await insertCodeCell('print("1")', { index: 1 });
+        await insertCodeCell('print("2")', { index: 2 });
+
+        // select second cell
+        // this tries to get the second cell selected by running it, but it doesn't work
+        // const secondCell = vscodeNotebook.activeNotebookEditor?.document.cells![1]!;
+        // await executeCell(secondCell);
+        // await waitForExecutionCompletedSuccessfully(secondCell);
+
+        // run command
+        await commandManager.executeCommand(
+            Commands.NativeNotebookRunAllCellsAbove,
+            vscodeNotebook.activeNotebookEditor?.document.uri!
+        );
+        const firstCell = vscodeNotebook.activeNotebookEditor?.document.cells![0]!;
+        await waitForExecutionCompletedSuccessfully(firstCell);
+
+        // The first cell should have a runState of Success
+        assert.strictEqual(firstCell?.metadata.runState, vscodeNotebookEnums.NotebookCellRunState.Success);
+    });
+});