Fixes to streamed output in native notebooks (#14158) (#155)

For #13611
* Fixes to streamed output
* added tests

src/client/datascience/jupyter/kernels/cellExecution.ts

@@ -5,8 +5,8 @@
 
 import { nbformat } from '@jupyterlab/coreutils';
 import type { KernelMessage } from '@jupyterlab/services/lib/kernel/messages';
-import { CancellationToken, CellOutputKind, CellStreamOutput, NotebookCell, NotebookCellRunState } from 'vscode';
-import type { NotebookEditor as VSCNotebookEditor } from '../../../../../types/vscode-proposed';
+import { CancellationToken, CellOutputKind, NotebookCell, NotebookCellRunState } from 'vscode';
+import type { CellDisplayOutput, NotebookEditor as VSCNotebookEditor } from '../../../../../types/vscode-proposed';
 import { concatMultilineString, formatStreamText } from '../../../../datascience-ui/common';
 import { IApplicationShell, IVSCodeNotebook } from '../../../common/application/types';
 import { traceInfo, traceWarning } from '../../../common/logger';
@@ -519,13 +519,17 @@ export class CellExecution {
             }
 
             // Might already have a stream message. If so, just add on to it.
+            // We use Rich output for text streams (not CellStreamOutput, known VSC Issues).
+            // https://github.com/microsoft/vscode-python/issues/14156
             const lastOutput =
                 exitingCellOutput.length > 0 ? exitingCellOutput[exitingCellOutput.length - 1] : undefined;
-            const existing: CellStreamOutput | undefined =
-                lastOutput && lastOutput.outputKind === CellOutputKind.Text ? lastOutput : undefined;
-            if (existing) {
+            const existing: CellDisplayOutput | undefined =
+                lastOutput && lastOutput.outputKind === CellOutputKind.Rich ? lastOutput : undefined;
+            if (existing && 'text/plain' in existing.data) {
                 // tslint:disable-next-line:restrict-plus-operands
-                existing.text = formatStreamText(concatMultilineString(existing.text + escape(msg.content.text)));
+                existing.data['text/plain'] = formatStreamText(
+                    concatMultilineString(`${existing.data['text/plain']}${escape(msg.content.text)}`)
+                );
                 edit.replaceCellOutput(this.cellIndex, [...exitingCellOutput]); // This is necessary to get VS code to update (for now)
             } else {
                 const originalText = formatStreamText(concatMultilineString(escape(msg.content.text)));

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

@@ -221,9 +221,7 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
         expect(displayCell.metadata.lastRunDuration).to.be.greaterThan(0, 'Duration should be > 0');
         expect(markdownOutput.data['text/markdown']).to.be.equal('foo', 'Display cell did not update');
     });
-    test('Clearing output while executing will ensure output is cleared', async function () {
-        // https://github.com/microsoft/vscode-python/issues/12302
-        return this.skip();
+    test('Clearing output while executing will ensure output is cleared', async () => {
         // Assume you are executing a cell that prints numbers 1-100.
         // When printing number 50, you click clear.
         // Cell output should now start printing output from 51 onwards, & not 1.
@@ -235,24 +233,35 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
                         time.sleep(0.1)
                         print(i)
 
-                    print("End")`
+                    print("End")`,
+            0
         );
         const cell = vscodeNotebook.activeNotebookEditor?.document.cells![0]!;
 
         await executeActiveDocument();
 
-        // Wait till execution count changes and status is error.
+        // Wait till we get the desired output.
         await waitForCondition(
-            async () => assertHasTextOutputInVSCode(cell, 'Start', 0, false),
+            async () =>
+                assertHasTextOutputInVSCode(cell, 'Start', 0, false) &&
+                assertHasTextOutputInVSCode(cell, '0', 0, false) &&
+                assertHasTextOutputInVSCode(cell, '1', 0, false) &&
+                assertHasTextOutputInVSCode(cell, '2', 0, false) &&
+                assertHasTextOutputInVSCode(cell, '3', 0, false) &&
+                assertHasTextOutputInVSCode(cell, '4', 0, false),
             15_000,
             'Cell did not get executed'
         );
 
         // Clear the cells
         await commands.executeCommand('notebook.clearAllCellsOutputs');
-        // Wait till execution count changes and status is error.
+
+        // Wait till previous output gets cleared & we have new output.
         await waitForCondition(
-            async () => assertNotHasTextOutputInVSCode(cell, 'Start', 0, false),
+            async () =>
+                assertNotHasTextOutputInVSCode(cell, 'Start', 0, false) &&
+                cell.outputs.length > 0 &&
+                cell.outputs[0].outputKind === vscodeNotebookEnums.CellOutputKind.Rich,
             5_000,
             'Cell did not get cleared'
         );
@@ -264,4 +273,74 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
         // Verify that it hasn't got added (even after interrupting).
         assertNotHasTextOutputInVSCode(cell, 'Start', 0, false);
     });
+    test('Clearing output via code', async () => {
+        // Assume you are executing a cell that prints numbers 1-100.
+        // When printing number 50, you click clear.
+        // Cell output should now start printing output from 51 onwards, & not 1.
+        await insertPythonCellAndWait(
+            dedent`
+                from IPython.display import display, clear_output
+                import time
+                print('foo')
+                display('foo')
+                time.sleep(2)
+                clear_output(True)
+                print('bar')
+                display('bar')`,
+            0
+        );
+        const cell = vscodeNotebook.activeNotebookEditor?.document.cells![0]!;
+
+        await executeActiveDocument();
+
+        // Wait for foo to be printed
+        await waitForCondition(
+            async () =>
+                assertHasTextOutputInVSCode(cell, 'foo', 0, false) &&
+                assertHasTextOutputInVSCode(cell, 'foo', 1, false),
+            15_000,
+            'Incorrect output'
+        );
+
+        // Wait for bar to be printed
+        await waitForCondition(
+            async () =>
+                assertHasTextOutputInVSCode(cell, 'bar', 0, false) &&
+                assertHasTextOutputInVSCode(cell, 'bar', 1, false),
+            15_000,
+            'Incorrect output'
+        );
+    });
+    test('Testing streamed output', async () => {
+        // Assume you are executing a cell that prints numbers 1-100.
+        // When printing number 50, you click clear.
+        // Cell output should now start printing output from 51 onwards, & not 1.
+        await insertPythonCellAndWait(
+            dedent`
+                    print("Start")
+                    import time
+                    for i in range(5):
+                        time.sleep(0.5)
+                        print(i)
+
+                    print("End")`,
+            0
+        );
+        const cell = vscodeNotebook.activeNotebookEditor?.document.cells![0]!;
+
+        await executeActiveDocument();
+
+        await waitForCondition(
+            async () =>
+                assertHasTextOutputInVSCode(cell, 'Start', 0, false) &&
+                assertHasTextOutputInVSCode(cell, '0', 0, false) &&
+                assertHasTextOutputInVSCode(cell, '1', 0, false) &&
+                assertHasTextOutputInVSCode(cell, '2', 0, false) &&
+                assertHasTextOutputInVSCode(cell, '3', 0, false) &&
+                assertHasTextOutputInVSCode(cell, '4', 0, false) &&
+                assertHasTextOutputInVSCode(cell, 'End', 0, false),
+            15_000,
+            'Incorrect output'
+        );
+    });
 });