Skip to content

Commit 03f84f7

Browse files
rchiodorchiodo
authored
Handle clear messages from jupyter (#5897)
* Handle clear messages from jupyter Fix a number of font/wrapping issues * Add news entry * Fix requirements
1 parent 6dff687 commit 03f84f7

7 files changed

Lines changed: 267 additions & 34 deletions

File tree

‎build/functional-test-requirements.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,3 +4,6 @@ jupyter
44
numpy
55
matplotlib
66
pandas
7+
torch
8+
scikit-learn
9+
livelossplot

‎news/2 Fixes/5801.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Add support for jupyter controls that clear.

‎src/client/datascience/jupyter/jupyterServer.ts‎

Lines changed: 69 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -634,24 +634,28 @@ export class JupyterServerBase implements INotebookServer {
634634
});
635635
}
636636

637+
const clearState : Map<string, boolean> = new Map<string, boolean>();
638+
637639
// Listen to the reponse messages and update state as we go
638640
if (request) {
639641
request.onIOPub = (msg: KernelMessage.IIOPubMessage) => {
640642
try {
641643
if (jupyterLab.KernelMessage.isExecuteResultMsg(msg)) {
642-
this.handleExecuteResult(msg as KernelMessage.IExecuteResultMsg, subscriber.cell);
644+
this.handleExecuteResult(msg as KernelMessage.IExecuteResultMsg, clearState, subscriber.cell);
643645
} else if (jupyterLab.KernelMessage.isExecuteInputMsg(msg)) {
644-
this.handleExecuteInput(msg as KernelMessage.IExecuteInputMsg, subscriber.cell);
646+
this.handleExecuteInput(msg as KernelMessage.IExecuteInputMsg, clearState, subscriber.cell);
645647
} else if (jupyterLab.KernelMessage.isStatusMsg(msg)) {
646-
this.handleStatusMessage(msg as KernelMessage.IStatusMsg, subscriber.cell);
648+
this.handleStatusMessage(msg as KernelMessage.IStatusMsg, clearState, subscriber.cell);
647649
} else if (jupyterLab.KernelMessage.isStreamMsg(msg)) {
648-
this.handleStreamMesssage(msg as KernelMessage.IStreamMsg, subscriber.cell);
650+
this.handleStreamMesssage(msg as KernelMessage.IStreamMsg, clearState, subscriber.cell);
649651
} else if (jupyterLab.KernelMessage.isDisplayDataMsg(msg)) {
650-
this.handleDisplayData(msg as KernelMessage.IDisplayDataMsg, subscriber.cell);
652+
this.handleDisplayData(msg as KernelMessage.IDisplayDataMsg, clearState, subscriber.cell);
651653
} else if (jupyterLab.KernelMessage.isUpdateDisplayDataMsg(msg)) {
652-
this.handleUpdateDisplayData(msg as KernelMessage.IUpdateDisplayDataMsg, subscriber.cell);
654+
this.handleUpdateDisplayData(msg as KernelMessage.IUpdateDisplayDataMsg, clearState, subscriber.cell);
655+
} else if (jupyterLab.KernelMessage.isClearOutputMsg(msg)) {
656+
this.handleClearOutput(msg as KernelMessage.IClearOutputMsg, clearState, subscriber.cell);
653657
} else if (jupyterLab.KernelMessage.isErrorMsg(msg)) {
654-
this.handleError(msg as KernelMessage.IErrorMsg, subscriber.cell);
658+
this.handleError(msg as KernelMessage.IErrorMsg, clearState, subscriber.cell);
655659
} else {
656660
this.logger.logWarning(`Unknown message ${msg.header.msg_type} : hasData=${'data' in msg.content}`);
657661
}
@@ -661,7 +665,7 @@ export class JupyterServerBase implements INotebookServer {
661665
subscriber.cell.data.execution_count = msg.content.execution_count as number;
662666
}
663667

664-
// Show our update if any new output
668+
// Show our update if any new output.
665669
subscriber.next(this.sessionStartTime);
666670
} catch (err) {
667671
// If not a restart error, then tell the subscriber
@@ -707,56 +711,80 @@ export class JupyterServerBase implements INotebookServer {
707711
});
708712
}
709713

710-
private addToCellData = (cell: ICell, output: nbformat.IUnrecognizedOutput | nbformat.IExecuteResult | nbformat.IDisplayData | nbformat.IStream | nbformat.IError) => {
711-
const data: nbformat.ICodeCell = cell.data as nbformat.ICodeCell;
712-
data.outputs = [...data.outputs, output];
713-
cell.data = data;
714+
private addToCellData = (cell: ICell, output: nbformat.IUnrecognizedOutput | nbformat.IExecuteResult | nbformat.IDisplayData | nbformat.IStream | nbformat.IError, clearState: Map<string, boolean>) => {
715+
// If a clear is pending, replace the output with the new one
716+
if (clearState.get(output.output_type)) {
717+
clearState.delete(output.output_type);
718+
const data: nbformat.ICodeCell = cell.data as nbformat.ICodeCell;
719+
const index = data.outputs.findIndex(o => o.output_type === output.output_type);
720+
if (index >= 0) {
721+
data.outputs.splice(index, 1, output);
722+
} else {
723+
data.outputs = [...data.outputs, output];
724+
}
725+
cell.data = data;
726+
} else {
727+
// Then append this data onto the end.
728+
const data: nbformat.ICodeCell = cell.data as nbformat.ICodeCell;
729+
data.outputs = [...data.outputs, output];
730+
cell.data = data;
731+
}
714732
}
715733

716-
private handleExecuteResult(msg: KernelMessage.IExecuteResultMsg, cell: ICell) {
717-
this.addToCellData(cell, { output_type: 'execute_result', data: msg.content.data, metadata: msg.content.metadata, execution_count: msg.content.execution_count });
734+
private handleExecuteResult(msg: KernelMessage.IExecuteResultMsg, clearState: Map<string, boolean>, cell: ICell) {
735+
this.addToCellData(
736+
cell,
737+
{ output_type: 'execute_result', data: msg.content.data, metadata: msg.content.metadata, execution_count: msg.content.execution_count },
738+
clearState);
718739
}
719740

720-
private handleExecuteInput(msg: KernelMessage.IExecuteInputMsg, cell: ICell) {
741+
private handleExecuteInput(msg: KernelMessage.IExecuteInputMsg, _clearState: Map<string, boolean>, cell: ICell) {
721742
cell.data.execution_count = msg.content.execution_count;
722743
}
723744

724-
private handleStatusMessage(msg: KernelMessage.IStatusMsg, cell: ICell) {
745+
private handleStatusMessage(msg: KernelMessage.IStatusMsg, _clearState: Map<string, boolean>, cell: ICell) {
725746
// Status change to idle generally means we finished. Not sure how to
726747
// make sure of this. Maybe only bother if an interrupt
727748
if (msg.content.execution_state === 'idle' && cell.state !== CellState.error) {
728749
cell.state = CellState.finished;
729750
}
730751
}
731752

732-
private handleStreamMesssage(msg: KernelMessage.IStreamMsg, cell: ICell) {
753+
private handleStreamMesssage(msg: KernelMessage.IStreamMsg, clearState: Map<string, boolean>, cell: ICell) {
733754
// Might already have a stream message. If so, just add on to it.
734755
const data: nbformat.ICodeCell = cell.data as nbformat.ICodeCell;
735756
const existing = data.outputs.find(o => o.output_type === 'stream');
736757
if (existing && existing.name === msg.content.name) {
737-
// tslint:disable-next-line:restrict-plus-operands
738-
existing.text = existing.text + msg.content.text;
758+
// If clear pending, then don't add.
759+
if (clearState.get('stream')) {
760+
clearState.delete('stream');
761+
existing.text = msg.content.text;
762+
} else {
763+
// tslint:disable-next-line:restrict-plus-operands
764+
existing.text = existing.text + msg.content.text;
765+
}
766+
739767
} else {
740768
// Create a new stream entry
741769
const output: nbformat.IStream = {
742770
output_type: 'stream',
743771
name: msg.content.name,
744772
text: msg.content.text
745773
};
746-
this.addToCellData(cell, output);
774+
this.addToCellData(cell, output, clearState);
747775
}
748776
}
749777

750-
private handleDisplayData(msg: KernelMessage.IDisplayDataMsg, cell: ICell) {
778+
private handleDisplayData(msg: KernelMessage.IDisplayDataMsg, clearState: Map<string, boolean>, cell: ICell) {
751779
const output: nbformat.IDisplayData = {
752780
output_type: 'display_data',
753781
data: msg.content.data,
754782
metadata: msg.content.metadata
755783
};
756-
this.addToCellData(cell, output);
784+
this.addToCellData(cell, output, clearState);
757785
}
758786

759-
private handleUpdateDisplayData(msg: KernelMessage.IUpdateDisplayDataMsg, cell: ICell) {
787+
private handleUpdateDisplayData(msg: KernelMessage.IUpdateDisplayDataMsg, _clearState: Map<string, boolean>, cell: ICell) {
760788
// Should already have a display data output in our cell.
761789
const data: nbformat.ICodeCell = cell.data as nbformat.ICodeCell;
762790
const output = data.outputs.find(o => o.output_type === 'display_data');
@@ -766,6 +794,21 @@ export class JupyterServerBase implements INotebookServer {
766794
}
767795
}
768796

797+
private handleClearOutput(msg: KernelMessage.IClearOutputMsg, clearState: Map<string, boolean>, cell: ICell) {
798+
// If the message says wait, add every message type to our clear state. This will
799+
// make us wait for this type of output before we clear it.
800+
if (msg && msg.content.wait) {
801+
clearState.set('display_data', true);
802+
clearState.set('error', true);
803+
clearState.set('execute_result', true);
804+
clearState.set('stream', true);
805+
} else {
806+
// Clear all outputs and start over again.
807+
const data: nbformat.ICodeCell = cell.data as nbformat.ICodeCell;
808+
data.outputs = [];
809+
}
810+
}
811+
769812
private handleInterrupted(cell: ICell) {
770813
this.handleError({
771814
channel: 'iopub',
@@ -781,17 +824,17 @@ export class JupyterServerBase implements INotebookServer {
781824
'KeyboardInterrupt: '
782825
]
783826
}
784-
}, cell);
827+
}, new Map<string, boolean>(), cell);
785828
}
786829

787-
private handleError(msg: KernelMessage.IErrorMsg, cell: ICell) {
830+
private handleError(msg: KernelMessage.IErrorMsg, clearState: Map<string, boolean>, cell: ICell) {
788831
const output: nbformat.IError = {
789832
output_type: 'error',
790833
ename: msg.content.ename,
791834
evalue: msg.content.evalue,
792835
traceback: msg.content.traceback
793836
};
794-
this.addToCellData(cell, output);
837+
this.addToCellData(cell, output, clearState);
795838
cell.state = CellState.error;
796839
}
797840
}

‎src/datascience-ui/history-react/cell.css‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -62,13 +62,21 @@
6262
.cell-output {
6363
margin-top: 5px;
6464
background: var(--override-widget-background, var(--vscode-notifications-background));
65-
white-space: pre-wrap;
66-
font-family: monospace;
6765
}
6866

69-
.cell-output pre {
67+
.cell-output-text {
68+
white-space: pre-wrap;
69+
font-size: var(--code-font-size);
70+
font-family: var(--code-font-family);
71+
}
72+
.cell-output-text pre {
7073
white-space: pre-wrap;
71-
font-family: monospace;
74+
font-size: var(--code-font-size);
75+
font-family: var(--code-font-family);
76+
}
77+
78+
.cell-output-html {
79+
white-space: none;
7280
}
7381

7482
.cell-output table {

‎src/datascience-ui/history-react/cell.tsx‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -293,7 +293,7 @@ export class Cell extends React.Component<ICellProps> {
293293
return [<Transform key={0} data={source}/>];
294294
}
295295

296-
private renderWithTransform = (mimetype: string, output : nbformat.IOutput, index : number, renderWithScrollbars: boolean, forceLightTheme: boolean) => {
296+
private renderWithTransform = (mimetype: string, output : nbformat.IOutput, index : number, renderWithScrollbars: boolean, forceLightTheme: boolean, isText: boolean) => {
297297

298298
// If we found a mimetype, use the transform
299299
if (mimetype) {
@@ -313,6 +313,7 @@ export class Cell extends React.Component<ICellProps> {
313313
if (mimetype === 'text/plain') {
314314
data = concatMultilineString(data as nbformat.MultilineString);
315315
renderWithScrollbars = true;
316+
isText = true;
316317
}
317318

318319
// Create a default set of properties
@@ -332,7 +333,9 @@ export class Cell extends React.Component<ICellProps> {
332333
style.color = this.invertColor(this.props.errorBackgroundColor);
333334
}
334335

335-
return <div id='stylewrapper' key={index} style={style}><Transform data={data} /></div>;
336+
const className = isText ? 'cell-output-text' : 'cell-output-html';
337+
338+
return <div id='stylewrapper' className={className} key={index} style={style}><Transform data={data} /></div>;
336339
}
337340
} catch (ex) {
338341
window.console.log('Error in rendering');
@@ -397,10 +400,13 @@ export class Cell extends React.Component<ICellProps> {
397400
// Only for text and error ouptut do we add scrollbars
398401
let addScrollbars = false;
399402
let forceLightTheme = false;
403+
let isText = false;
400404

401405
// Stream and error output need to be converted
402406
if (copy.output_type === 'stream') {
403407
addScrollbars = true;
408+
isText = true;
409+
404410
// Stream output needs to be wrapped in xmp so it
405411
// show literally. Otherwise < chars start a new html element.
406412
const stream = copy as nbformat.IStream;
@@ -427,6 +433,7 @@ export class Cell extends React.Component<ICellProps> {
427433
} else if (copy.output_type === 'error') {
428434
addScrollbars = true;
429435
forceLightTheme = true;
436+
isText = true;
430437
const error = copy as nbformat.IError;
431438
try {
432439
const converter = new ansiToHtml();
@@ -450,7 +457,7 @@ export class Cell extends React.Component<ICellProps> {
450457

451458
// If that worked, use the transform
452459
if (mimetype) {
453-
return this.renderWithTransform(mimetype, copy, index, addScrollbars, forceLightTheme);
460+
return this.renderWithTransform(mimetype, copy, index, addScrollbars, forceLightTheme, isText);
454461
}
455462

456463
if (copy.data) {

‎src/test/datascience/history.functional.test.tsx‎

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,22 +2,26 @@
22
// Licensed under the MIT License.
33
'use strict';
44
import * as assert from 'assert';
5+
import * as fs from 'fs-extra';
6+
import { parse } from 'node-html-parser';
57
import * as path from 'path';
68
import * as TypeMoq from 'typemoq';
79
import { Disposable, TextDocument, TextEditor } from 'vscode';
810

911
import { IApplicationShell, IDocumentManager } from '../../client/common/application/types';
1012
import { createDeferred } from '../../client/common/utils/async';
1113
import { noop } from '../../client/common/utils/misc';
14+
import { generateCellsFromDocument } from '../../client/datascience/cellFactory';
15+
import { concatMultilineString } from '../../client/datascience/common';
1216
import { EditorContexts } from '../../client/datascience/constants';
1317
import { HistoryMessageListener } from '../../client/datascience/history/historyMessageListener';
1418
import { HistoryMessages } from '../../client/datascience/history/historyTypes';
1519
import { IHistory, IHistoryProvider } from '../../client/datascience/types';
1620
import { CellButton } from '../../datascience-ui/history-react/cellButton';
1721
import { MainPanel } from '../../datascience-ui/history-react/MainPanel';
18-
//import { asyncDump } from '../common/asyncDump';
1922
import { sleep } from '../core';
2023
import { DataScienceIocContainer } from './dataScienceIocContainer';
24+
import { createDocument } from './editor-integration/helpers';
2125
import {
2226
addCode,
2327
addContinuousMockData,
@@ -40,6 +44,7 @@ import {
4044
} from './historyTestHelpers';
4145
import { waitForUpdate } from './reactHelpers';
4246

47+
//import { asyncDump } from '../common/asyncDump';
4348
// tslint:disable:max-func-body-length trailing-comma no-any no-multiline-string
4449
suite('DataScience History output tests', () => {
4550
const disposables: Disposable[] = [];
@@ -539,4 +544,38 @@ for _ in range(50):
539544

540545
verifyHtmlOnCell(wrapper, '<img', CellPosition.Last);
541546
}, () => { return ioc; });
547+
548+
runMountedTest('LiveLossPlot', async (wrapper) => {
549+
// Only run this test when not mocking. Too complicated to mimic otherwise
550+
if (!ioc.mockJupyter) {
551+
// Load all of our cells
552+
const testFile = path.join(srcDirectory(), 'liveloss.py');
553+
const version = 1;
554+
const inputText = await fs.readFile(testFile, 'utf-8');
555+
const document = createDocument(inputText, testFile, version, TypeMoq.Times.atLeastOnce(), true);
556+
const cells = generateCellsFromDocument(document.object);
557+
assert.ok(cells, 'No cells generated');
558+
assert.equal(cells.length, 8, 'Not enough cells generated');
559+
560+
// Run the first 7 cells
561+
for (let i = 0; i < 7; i += 1) {
562+
const renderCount = i === 2 ? 5 : 4 ; // Cell 2 outputs a print statement
563+
await addCode(getOrCreateHistory, wrapper, concatMultilineString(cells[i].data.source), renderCount);
564+
}
565+
566+
// Last cell should generate a series of updates. Verify we end up with a single image
567+
await addCode(getOrCreateHistory, wrapper, concatMultilineString(cells[7].data.source), 25);
568+
const cell = getLastOutputCell(wrapper);
569+
570+
const output = cell!.find('div.cell-output');
571+
assert.ok(output.length > 0, 'No output cell found');
572+
const outHtml = output.html();
573+
574+
const root = parse(outHtml) as any;
575+
const imgs = root.querySelectorAll('img') as HTMLElement[];
576+
assert.ok(imgs, 'No images found');
577+
assert.equal(imgs.length, 1, 'Wrong number of images');
578+
}
579+
580+
}, () => { return ioc; });
542581
});

0 commit comments

Comments
 (0)