Remove refBool with better comments

src/notebooks/execution/cellExecutionMessageHandler.ts

@@ -30,7 +30,6 @@ import { CellExecutionCreator } from './cellExecutionCreator';
 import { IApplicationShell } from '../../platform/common/application/types';
 import { disposeAllDisposables } from '../../platform/common/helpers';
 import { traceError, traceWarning } from '../../platform/logging';
-import { RefBool } from '../../platform/common/refBool';
 import { IDisposable, IExtensionContext } from '../../platform/common/types';
 import { traceCellMessage, cellOutputToVSCCellOutput, translateCellDisplayOutput, isJupyterNotebook } from '../helpers';
 import { formatStreamText, concatMultilineString } from '../../webviews/webview-side/common';
@@ -73,10 +72,12 @@ export class CellExecutionMessageHandler implements IDisposable {
      */
     private completedExecution?: boolean;
     /**
-     * Listen to messages and update our cell execution state appropriately
-     * Keep track of our clear state
+     * Jupyter can sent a `clear_output` message which indicates the output of a cell should be cleared.
+     * If the flag `wait` is set to `true`, then we should wait for the next output before clearing the output.
+     * I.e. if the value for `wait` is false (default) then clear the cell output immediately.
+     * https://ipywidgets.readthedocs.io/en/latest/examples/Output%20Widget.html#Output-widgets:-leveraging-Jupyter's-display-system
      */
-    private readonly clearState = new RefBool(false);
+    private clearOutputOnNextUpdateToOutput?: boolean;
 
     private execution?: NotebookCellExecution;
     private readonly _onErrorHandlingIOPubMessage = new EventEmitter<{
@@ -333,7 +334,18 @@ export class CellExecutionMessageHandler implements IDisposable {
             this.execution.executionOrder = msg.content.execution_count;
         }
     }
-
+    private clearOutputIfNecessary(execution: NotebookCellExecution | undefined): {
+        previousValueOfClearOutputOnNextUpdateToOutput: boolean;
+    } {
+        if (this.clearOutputOnNextUpdateToOutput) {
+            traceCellMessage(this.cell, 'Clear cell output');
+            this.clearLastUsedStreamOutput();
+            execution?.clearOutput().then(noop, noop);
+            this.clearOutputOnNextUpdateToOutput = false;
+            return { previousValueOfClearOutputOnNextUpdateToOutput: true };
+        }
+        return { previousValueOfClearOutputOnNextUpdateToOutput: false };
+    }
     private addToCellData(output: ExecuteResult | DisplayData | nbformat.IStream | nbformat.IError | nbformat.IOutput) {
         if (
             this.context.extensionMode === ExtensionMode.Test &&
@@ -358,11 +370,7 @@ export class CellExecutionMessageHandler implements IDisposable {
         }
         traceCellMessage(this.cell, 'Update output');
         // Clear if necessary
-        if (this.clearState.value) {
-            this.clearLastUsedStreamOutput();
-            this.execution?.clearOutput().then(noop, noop);
-            this.clearState.update(false);
-        }
+        this.clearOutputIfNecessary(this.execution);
         // Keep track of the displa_id against the output item, we might need this to update this later.
         if (displayId) {
             this.outputDisplayIdTracker.trackOutputByDisplayId(this.cell, displayId, cellOutput);
@@ -485,13 +493,7 @@ export class CellExecutionMessageHandler implements IDisposable {
         const task = this.execution || this.createTemporaryTask();
 
         // Clear output if waiting for a clear
-        const clearOutput = this.clearState.value;
-        if (clearOutput) {
-            traceCellMessage(this.cell, 'Clear cell output');
-            this.clearLastUsedStreamOutput();
-            task?.clearOutput().then(noop, noop);
-            this.clearState.update(false);
-        }
+        const { previousValueOfClearOutputOnNextUpdateToOutput } = this.clearOutputIfNecessary(task);
         // Ensure we append to previous output, only if the streams as the same &
         // If the last output is the desired stream type.
         if (this.lastUsedStreamOutput?.stream === msg.content.name) {
@@ -524,7 +526,7 @@ export class CellExecutionMessageHandler implements IDisposable {
             });
             traceCellMessage(this.cell, `Replace output items ${this.lastUsedStreamOutput.text.substring(0, 100)}`);
             task?.replaceOutputItems(output.items, this.lastUsedStreamOutput.output).then(noop, noop);
-        } else if (clearOutput) {
+        } else if (previousValueOfClearOutputOnNextUpdateToOutput) {
             // Replace the current outputs with a single new output.
             const text = formatStreamText(concatMultilineString(msg.content.text));
             const output = cellOutputToVSCCellOutput({
@@ -565,7 +567,7 @@ export class CellExecutionMessageHandler implements IDisposable {
         // If the message says wait, add every message type to our clear state. This will
         // make us wait for this type of output before we clear it.
         if (msg && msg.content.wait) {
-            this.clearState.update(true);
+            this.clearOutputOnNextUpdateToOutput = true;
         } else {
             // Possible execution of cell has completed (the task would have been disposed).
             // This message could have come from a background thread.

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

@@ -371,6 +371,66 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
 
         await waitForExecutionCompletedSuccessfully(cell);
     });
+    test('Clearing output immediately via code', async function () {
+        // 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 insertCodeCell(
+            dedent`
+            from ipywidgets import widgets
+            from IPython.display import display, clear_output
+            import time
+
+            display(widgets.Button(description="First Button"))
+
+            time.sleep(2)
+            clear_output()
+
+            display(widgets.Button(description="Second Button"))`,
+            { index: 0 }
+        );
+        const cell = vscodeNotebook.activeNotebookEditor?.notebook.cellAt(0)!;
+
+        await runAllCellsInActiveNotebook();
+
+        await Promise.all([
+            waitForExecutionCompletedSuccessfully(cell),
+            waitForTextOutput(cell, 'Second Button', 1, false)
+        ]);
+    });
+    test('Clearing output via code only when receiving new output', async function () {
+        // 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 insertCodeCell(
+            dedent`
+            from ipywidgets import widgets
+            from IPython.display import display, clear_output
+            import time
+
+            display(widgets.Button(description="First Button"))
+
+            time.sleep(2)
+            clear_output(True)
+
+            display(widgets.Button(description="Second Button"))`,
+            { index: 0 }
+        );
+        const cell = vscodeNotebook.activeNotebookEditor?.notebook.cellAt(0)!;
+
+        await runAllCellsInActiveNotebook();
+
+        // Wait for first button to appear then second.
+        await Promise.all([
+            waitForExecutionCompletedSuccessfully(cell),
+            waitForTextOutput(cell, 'First Button', 0, false),
+            waitForTextOutput(cell, 'Second Button', 1, false)
+        ]);
+
+        // Verify first is no longer visible and second is visible.
+        assert.notInclude(getCellOutputs(cell), 'First Button');
+        assert.include(getCellOutputs(cell), 'Second Button');
+    });
     test('Shell commands should give preference to executables in Python Environment', async function () {
         if (IS_REMOTE_NATIVE_TEST()) {
             return this.skip();