Merge pull request #10155 from microsoft/kernelConnectionWrapperStep2

src/notebooks/execution/cellExecutionMessageHandlerService.ts

@@ -0,0 +1,71 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { Kernel, KernelMessage } from '@jupyterlab/services';
+import { NotebookCell, NotebookCellExecution, NotebookController, NotebookDocument, workspace } from 'vscode';
+import { ITracebackFormatter } from '../../kernels/types';
+import { IApplicationShell } from '../../platform/common/application/types';
+import { disposeAllDisposables } from '../../platform/common/helpers';
+import { IDisposable, IExtensionContext } from '../../platform/common/types';
+import { CellOutputDisplayIdTracker } from './cellDisplayIdTracker';
+import { CellExecutionMessageHandler } from './cellExecutionMessageHandler';
+
+export class CellExecutionMessageHandlerService {
+    private readonly disposables: IDisposable[] = [];
+    private notebook?: NotebookDocument;
+    private readonly messageHandlers = new WeakMap<NotebookCell, CellExecutionMessageHandler>();
+    constructor(
+        private readonly appShell: IApplicationShell,
+        private readonly controller: NotebookController,
+        private readonly outputDisplayIdTracker: CellOutputDisplayIdTracker,
+        private readonly context: IExtensionContext,
+        private readonly formatters: ITracebackFormatter[]
+    ) {
+        workspace.onDidChangeNotebookDocument(
+            (e) => {
+                if (e.notebook !== this.notebook) {
+                    return;
+                }
+                e.contentChanges.forEach((change) =>
+                    // If the cell is deleted, then dispose the corresponding handler.
+                    change.removedCells.forEach((cell) => this.messageHandlers.get(cell)?.dispose())
+                );
+            },
+            this,
+            this.disposables
+        );
+    }
+    dispose() {
+        disposeAllDisposables(this.disposables);
+        if (this.notebook) {
+            this.notebook.getCells().forEach((cell) => this.messageHandlers.get(cell)?.dispose());
+        }
+    }
+    public registerListener(
+        cell: NotebookCell,
+        options: {
+            kernel: Kernel.IKernelConnection;
+            request: Kernel.IShellFuture<KernelMessage.IExecuteRequestMsg, KernelMessage.IExecuteReplyMsg>;
+            cellExecution: NotebookCellExecution;
+            onErrorHandlingExecuteRequestIOPubMessage: (error: Error) => void;
+        }
+    ): CellExecutionMessageHandler {
+        this.notebook = cell.notebook;
+        // Always dispose any previous handlers & create new ones.
+        this.messageHandlers.get(cell)?.dispose();
+        const handler = new CellExecutionMessageHandler(
+            cell,
+            this.appShell,
+            this.controller,
+            this.outputDisplayIdTracker,
+            this.context,
+            this.formatters,
+            options.kernel,
+            options.request,
+            options.cellExecution
+        );
+        // This object must be kept in memory has it monitors the kernel messages.
+        this.messageHandlers.set(cell, handler);
+        return handler;
+    }
+}

src/notebooks/execution/kernelExecution.ts

@@ -27,6 +27,7 @@ import {
 import { traceCellMessage } from '../helpers';
 import { getDisplayPath } from '../../platform/common/platform/fs-paths';
 import { getAssociatedNotebookDocument } from '../controllers/kernelSelector';
+import { CellExecutionMessageHandlerService } from './cellExecutionMessageHandlerService';
 
 /**
  * Separate class that deals just with kernel execution.
@@ -49,7 +50,15 @@ export class KernelExecution implements IDisposable {
         context: IExtensionContext,
         formatters: ITracebackFormatter[]
     ) {
-        this.executionFactory = new CellExecutionFactory(appShell, controller, outputTracker, context, formatters);
+        const requestListener = new CellExecutionMessageHandlerService(
+            appShell,
+            controller,
+            outputTracker,
+            context,
+            formatters
+        );
+        this.disposables.push(requestListener);
+        this.executionFactory = new CellExecutionFactory(controller, requestListener);
     }
 
     public get onPreExecute() {

src/platform/common/refBool.node.ts → src/platform/common/refBool.ts

@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 export class RefBool {
     constructor(private val: boolean) {}
 

src/test/datascience/notebook/helper.ts

@@ -830,7 +830,10 @@ export async function waitForTextOutput(
                 cell.index + 1
             } in output index ${index}, it is ${cell.outputs
                 .map(
-                    (output, index) => `Output for Index "${index}" is "${output.items.map(getOutputText).join('\n')}"`
+                    (output, index) =>
+                        `Output for Index "${index}" with total outputs ${output.items.length} is "${output.items
+                            .map(getOutputText)
+                            .join('\n')}"`
                 )
                 .join('\n')}`
     );

src/test/datascience/widgets/notebooks/button_widget_comm_msg.ipynb

@@ -0,0 +1,39 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import ipywidgets as widgets\n",
+    "from IPython.display import display\n",
+    "button = widgets.Button(description=\"Click Me!\")\n",
+    "\n",
+    "display(button)\n",
+    "\n",
+    "def on_button_clicked(b):\n",
+    "    print(\"Button clicked.\")\n",
+    "\n",
+    "button.on_click(on_button_clicked)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "display(button)"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

src/test/datascience/widgets/notebooks/button_widget_comm_msg_matplotlib.ipynb

@@ -0,0 +1,47 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib widget\n",
+    "from IPython.display import display\n",
+    "from ipywidgets import widgets\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "\n",
+    "\n",
+    "button = widgets.Button(description=\"Click Me!\")\n",
+    "\n",
+    "def on_button_clicked(b):\n",
+    "\tX = np.linspace(0, 2*np.pi)\n",
+    "\tY = np.sin(X)\n",
+    "\n",
+    "\t_, ax = plt.subplots()\n",
+    "\tax.plot(X, Y)\n",
+    "\n",
+    "display(button)\n",
+    "button.on_click(on_button_clicked)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "display(button)"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

src/test/datascience/widgets/standard.vscode.common.ts

@@ -24,6 +24,7 @@ import {
     runCell,
     waitForExecutionCompletedSuccessfully,
     waitForKernelToGetAutoSelected,
+    waitForTextOutput,
     workAroundVSCodeNotebookStartPages
 } from '../notebook/helper';
 import { initializeWidgetComms, Utils } from './commUtils';
@@ -214,6 +215,30 @@ export function sharedIPyWidgetsTests(
         await assertOutputContainsHtml(cell1, comms, ['Button clicked']);
         await assertOutputContainsHtml(cell2, comms, ['Button clicked']);
     });
+    test('Button Widget with custom comm message', async () => {
+        const comms = await initializeNotebook({ templateFile: 'button_widget_comm_msg.ipynb' });
+        const [cell0, cell1] = vscodeNotebook.activeNotebookEditor!.notebook.getCells();
+
+        await executeCellAndWaitForOutput(cell0, comms);
+        await executeCellAndWaitForOutput(cell1, comms);
+        await assertOutputContainsHtml(cell0, comms, ['Click Me!', '<button']);
+
+        // Click the button and verify we have output in the same cell.
+        await click(comms, cell0, 'button');
+        await waitForTextOutput(cell0, 'Button clicked.', 1, false);
+    });
+    test('Button Widget with custom comm message rendering a matplotlib widget', async () => {
+        const comms = await initializeNotebook({ templateFile: 'button_widget_comm_msg_matplotlib.ipynb' });
+        const [cell0, cell1] = vscodeNotebook.activeNotebookEditor!.notebook.getCells();
+
+        await executeCellAndWaitForOutput(cell0, comms);
+        await executeCellAndWaitForOutput(cell1, comms);
+        await assertOutputContainsHtml(cell0, comms, ['Click Me!', '<button']);
+
+        // Click the button and verify we have output in the same cell.
+        await click(comms, cell0, 'button');
+        await assertOutputContainsHtml(cell0, comms, ['>Figure 1<', '<canvas', 'Download plot']);
+    });
     test('Render IPySheets', async () => {
         const comms = await initializeNotebook({ templateFile: 'ipySheet_widgets.ipynb' });
         const [, cell1] = vscodeNotebook.activeNotebookEditor!.notebook.getCells();