Skip to content

Commit 4230c22

Browse files
DonJayamanneDonJayamanne
authored
Compress notebook output streams before rendering (microsoft#160667)
* Compress notebook output streams before rendering * OOps * Combine the buffers manually * Address code review * oops * Fixes * We can have multiple stream mimes in an output * oops
1 parent 1d500fb commit 4230c22

File tree

6 files changed

+149
-21
lines changed

6 files changed

+149
-21
lines changed

‎extensions/ipynb/src/serializers.ts

+12-15
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import type * as nbformat from '@jupyterlab/nbformat';
77
import { NotebookCell, NotebookCellData, NotebookCellKind, NotebookCellOutput } from 'vscode';
88
import { CellMetadata, CellOutputMetadata } from './common';
99
import { textMimeTypes } from './deserializers';
10+
import { compressOutputItemStreams } from './streamCompressor';
1011

1112
const textDecoder = new TextDecoder();
1213

@@ -270,21 +271,17 @@ type JupyterOutput =
270271

271272
function convertStreamOutput(output: NotebookCellOutput): JupyterOutput {
272273
const outputs: string[] = [];
273-
output.items
274-
.filter((opit) => opit.mime === CellOutputMimeTypes.stderr || opit.mime === CellOutputMimeTypes.stdout)
275-
.map((opit) => textDecoder.decode(opit.data))
276-
.forEach(value => {
277-
// Ensure each line is a seprate entry in an array (ending with \n).
278-
const lines = value.split('\n');
279-
// If the last item in `outputs` is not empty and the first item in `lines` is not empty, then concate them.
280-
// As they are part of the same line.
281-
if (outputs.length && lines.length && lines[0].length > 0) {
282-
outputs[outputs.length - 1] = `${outputs[outputs.length - 1]}${lines.shift()!}`;
283-
}
284-
for (const line of lines) {
285-
outputs.push(line);
286-
}
287-
});
274+
const compressedStream = output.items.length ? new TextDecoder().decode(compressOutputItemStreams(output.items[0].mime, output.items)) : '';
275+
// Ensure each line is a separate entry in an array (ending with \n).
276+
const lines = compressedStream.split('\n');
277+
// If the last item in `outputs` is not empty and the first item in `lines` is not empty, then concate them.
278+
// As they are part of the same line.
279+
if (outputs.length && lines.length && lines[0].length > 0) {
280+
outputs[outputs.length - 1] = `${outputs[outputs.length - 1]}${lines.shift()!}`;
281+
}
282+
for (const line of lines) {
283+
outputs.push(line);
284+
}
288285

289286
for (let index = 0; index < (outputs.length - 1); index++) {
290287
outputs[index] = `${outputs[index]}\n`;

‎extensions/ipynb/src/streamCompressor.ts

+63
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License. See License.txt in the project root for license information.
4+
*--------------------------------------------------------------------------------------------*/
5+
6+
import type { NotebookCellOutputItem } from 'vscode';
7+
8+
9+
/**
10+
* Given a stream of individual stdout outputs, this function will return the compressed lines, escaping some of the common terminal escape codes.
11+
* E.g. some terminal escape codes would result in the previous line getting cleared, such if we had 3 lines and
12+
* last line contained such a code, then the result string would be just the first two lines.
13+
*/
14+
export function compressOutputItemStreams(mimeType: string, outputs: NotebookCellOutputItem[]) {
15+
// return outputs.find(op => op.mime === mimeType)!.data.buffer;
16+
17+
const buffers: Uint8Array[] = [];
18+
let startAppending = false;
19+
// Pick the first set of outputs with the same mime type.
20+
for (const output of outputs) {
21+
if (output.mime === mimeType) {
22+
if ((buffers.length === 0 || startAppending)) {
23+
buffers.push(output.data);
24+
startAppending = true;
25+
}
26+
} else if (startAppending) {
27+
startAppending = false;
28+
}
29+
}
30+
compressStreamBuffer(buffers);
31+
const totalBytes = buffers.reduce((p, c) => p + c.byteLength, 0);
32+
const combinedBuffer = new Uint8Array(totalBytes);
33+
let offset = 0;
34+
for (const buffer of buffers) {
35+
combinedBuffer.set(buffer, offset);
36+
offset = offset + buffer.byteLength;
37+
}
38+
return combinedBuffer;
39+
}
40+
const MOVE_CURSOR_1_LINE_COMMAND = `${String.fromCharCode(27)}[A`;
41+
const MOVE_CURSOR_1_LINE_COMMAND_BYTES = MOVE_CURSOR_1_LINE_COMMAND.split('').map(c => c.charCodeAt(0));
42+
const LINE_FEED = 10;
43+
function compressStreamBuffer(streams: Uint8Array[]) {
44+
streams.forEach((stream, index) => {
45+
if (index === 0 || stream.length < MOVE_CURSOR_1_LINE_COMMAND.length) {
46+
return;
47+
}
48+
49+
const previousStream = streams[index - 1];
50+
51+
// Remove the previous line if required.
52+
const command = stream.subarray(0, MOVE_CURSOR_1_LINE_COMMAND.length);
53+
if (command[0] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[0] && command[1] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[1] && command[2] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[2]) {
54+
const lastIndexOfLineFeed = previousStream.lastIndexOf(LINE_FEED);
55+
if (lastIndexOfLineFeed === -1) {
56+
return;
57+
}
58+
streams[index - 1] = previousStream.subarray(0, lastIndexOfLineFeed);
59+
streams[index] = stream.subarray(MOVE_CURSOR_1_LINE_COMMAND.length);
60+
}
61+
});
62+
return streams;
63+
}

‎src/vs/workbench/api/common/extHostTypes.ts

+3-2
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ import { IExtensionDescription } from 'vs/platform/extensions/common/extensions'
1717
import { FileSystemProviderErrorCode, markAsFileSystemProviderError } from 'vs/platform/files/common/files';
1818
import { RemoteAuthorityResolverErrorCode } from 'vs/platform/remote/common/remoteAuthorityResolver';
1919
import { IRelativePatternDto } from 'vs/workbench/api/common/extHost.protocol';
20-
import { CellEditType, ICellPartialMetadataEdit, IDocumentMetadataEdit } from 'vs/workbench/contrib/notebook/common/notebookCommon';
20+
import { CellEditType, ICellPartialMetadataEdit, IDocumentMetadataEdit, isTextStreamMime } from 'vs/workbench/contrib/notebook/common/notebookCommon';
2121
import { checkProposedApiEnabled } from 'vs/workbench/services/extensions/common/extensions';
2222
import type * as vscode from 'vscode';
2323

@@ -3517,7 +3517,8 @@ export class NotebookCellOutput {
35173517
for (let i = 0; i < items.length; i++) {
35183518
const item = items[i];
35193519
const normalMime = normalizeMimeType(item.mime);
3520-
if (!seen.has(normalMime)) {
3520+
// We can have multiple text stream mime types in the same output.
3521+
if (!seen.has(normalMime) || isTextStreamMime(normalMime)) {
35213522
seen.add(normalMime);
35223523
continue;
35233524
}

‎src/vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView.ts

+7-4
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ import { NOTEBOOK_WEBVIEW_BOUNDARY } from 'vs/workbench/contrib/notebook/browser
3737
import { preloadsScriptStr } from 'vs/workbench/contrib/notebook/browser/view/renderers/webviewPreloads';
3838
import { transformWebviewThemeVars } from 'vs/workbench/contrib/notebook/browser/view/renderers/webviewThemeMapping';
3939
import { MarkupCellViewModel } from 'vs/workbench/contrib/notebook/browser/viewModel/markupCellViewModel';
40-
import { CellUri, INotebookRendererInfo, NotebookSetting, RendererMessagingSpec } from 'vs/workbench/contrib/notebook/common/notebookCommon';
40+
import { CellUri, INotebookRendererInfo, isTextStreamMime, NotebookSetting, RendererMessagingSpec } from 'vs/workbench/contrib/notebook/common/notebookCommon';
4141
import { INotebookKernel } from 'vs/workbench/contrib/notebook/common/notebookKernelService';
4242
import { IScopedRendererMessaging } from 'vs/workbench/contrib/notebook/common/notebookRendererMessagingService';
4343
import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookService';
@@ -46,6 +46,7 @@ import { WebviewWindowDragMonitor } from 'vs/workbench/contrib/webview/browser/w
4646
import { IEditorGroupsService } from 'vs/workbench/services/editor/common/editorGroupsService';
4747
import { IWorkbenchEnvironmentService } from 'vs/workbench/services/environment/common/environmentService';
4848
import { FromWebviewMessage, IAckOutputHeight, IClickedDataUrlMessage, ICodeBlockHighlightRequest, IContentWidgetTopRequest, IControllerPreload, ICreationContent, ICreationRequestMessage, IFindMatch, IMarkupCellInitialization, RendererMetadata, ToWebviewMessage } from './webviewMessages';
49+
import { compressOutputItemStreams } from 'vs/workbench/contrib/notebook/browser/view/renderers/stdOutErrorPreProcessor';
4950

5051
export interface ICachedInset<K extends ICommonCellInfo> {
5152
outputId: string;
@@ -1278,12 +1279,14 @@ var requirejs = (function() {
12781279
let updatedContent: ICreationContent | undefined = undefined;
12791280
if (content.type === RenderOutputType.Extension) {
12801281
const output = content.source.model;
1281-
const first = output.outputs.find(op => op.mime === content.mimeType)!;
1282+
const firstBuffer = isTextStreamMime(content.mimeType) ?
1283+
compressOutputItemStreams(content.mimeType, output.outputs) :
1284+
output.outputs.find(op => op.mime === content.mimeType)!.data.buffer;
12821285
updatedContent = {
12831286
type: RenderOutputType.Extension,
12841287
outputId: outputCache.outputId,
1285-
mimeType: first.mime,
1286-
valueBytes: first.data.buffer,
1288+
mimeType: content.mimeType,
1289+
valueBytes: firstBuffer,
12871290
metadata: output.metadata,
12881291
};
12891292
}

‎src/vs/workbench/contrib/notebook/browser/view/renderers/stdOutErrorPreProcessor.ts

+56
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License. See License.txt in the project root for license information.
4+
*--------------------------------------------------------------------------------------------*/
5+
6+
import { VSBuffer } from 'vs/base/common/buffer';
7+
import type { IOutputItemDto } from 'vs/workbench/contrib/notebook/common/notebookCommon';
8+
9+
10+
/**
11+
* Given a stream of individual stdout outputs, this function will return the compressed lines, escaping some of the common terminal escape codes.
12+
* E.g. some terminal escape codes would result in the previous line getting cleared, such if we had 3 lines and
13+
* last line contained such a code, then the result string would be just the first two lines.
14+
*/
15+
export function compressOutputItemStreams(mimeType: string, outputs: IOutputItemDto[]) {
16+
const buffers: Uint8Array[] = [];
17+
let startAppending = false;
18+
19+
// Pick the first set of outputs with the same mime type.
20+
for (const output of outputs) {
21+
if (output.mime === mimeType) {
22+
if ((buffers.length === 0 || startAppending)) {
23+
buffers.push(output.data.buffer);
24+
startAppending = true;
25+
}
26+
} else if (startAppending) {
27+
startAppending = false;
28+
}
29+
}
30+
compressStreamBuffer(buffers);
31+
return VSBuffer.concat(buffers.map(buffer => VSBuffer.wrap(buffer))).buffer;
32+
}
33+
const MOVE_CURSOR_1_LINE_COMMAND = `${String.fromCharCode(27)}[A`;
34+
const MOVE_CURSOR_1_LINE_COMMAND_BYTES = MOVE_CURSOR_1_LINE_COMMAND.split('').map(c => c.charCodeAt(0));
35+
const LINE_FEED = 10;
36+
function compressStreamBuffer(streams: Uint8Array[]) {
37+
streams.forEach((stream, index) => {
38+
if (index === 0 || stream.length < MOVE_CURSOR_1_LINE_COMMAND.length) {
39+
return;
40+
}
41+
42+
const previousStream = streams[index - 1];
43+
44+
// Remove the previous line if required.
45+
const command = stream.subarray(0, MOVE_CURSOR_1_LINE_COMMAND.length);
46+
if (command[0] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[0] && command[1] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[1] && command[2] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[2]) {
47+
const lastIndexOfLineFeed = previousStream.lastIndexOf(LINE_FEED);
48+
if (lastIndexOfLineFeed === -1) {
49+
return;
50+
}
51+
streams[index - 1] = previousStream.subarray(0, lastIndexOfLineFeed);
52+
streams[index] = stream.subarray(MOVE_CURSOR_1_LINE_COMMAND.length);
53+
}
54+
});
55+
return streams;
56+
}

‎src/vs/workbench/contrib/notebook/common/notebookCommon.ts

+8
Original file line numberDiff line numberDiff line change
@@ -947,3 +947,11 @@ export interface NotebookExtensionDescription {
947947
readonly id: ExtensionIdentifier;
948948
readonly location: UriComponents | undefined;
949949
}
950+
951+
/**
952+
* Whether the provided mime type is a text streamn like `stdout`, `stderr`.
953+
*/
954+
export function isTextStreamMime(mimeType: string) {
955+
return ['application/vnd.code.notebook.stdout', 'application/x.notebook.stdout', 'application/x.notebook.stream', 'application/vnd.code.notebook.stderr', 'application/x.notebook.stderr'].includes(mimeType);
956+
}
957+

0 commit comments

Comments
 (0)