Compress streams in notebook outputs (#160946)

* Revert "Compress notebook output streams before rendering (#160667)"

This reverts commit 4230c22a08.

* Compress stream output items

* Minor perf improvements

* Misc

* Comments

* Added tests

* Merge issue

* More merge issues

* Misc

* Address code review comments
This commit is contained in:
Don Jayamanne
2022-10-12 08:43:46 +11:00
committed by GitHub
parent 4322170fd8
commit 43957ccfe1
9 changed files with 341 additions and 141 deletions

View File

@@ -7,7 +7,6 @@ import type * as nbformat from '@jupyterlab/nbformat';
import { NotebookCell, NotebookCellData, NotebookCellKind, NotebookCellOutput } from 'vscode';
import { CellOutputMetadata } from './common';
import { textMimeTypes } from './deserializers';
import { compressOutputItemStreams } from './streamCompressor';
const textDecoder = new TextDecoder();
@@ -277,17 +276,21 @@ type JupyterOutput =
function convertStreamOutput(output: NotebookCellOutput): JupyterOutput {
const outputs: string[] = [];
const compressedStream = output.items.length ? new TextDecoder().decode(compressOutputItemStreams(output.items[0].mime, output.items)) : '';
// Ensure each line is a separate entry in an array (ending with \n).
const lines = compressedStream.split('\n');
// If the last item in `outputs` is not empty and the first item in `lines` is not empty, then concate them.
// As they are part of the same line.
if (outputs.length && lines.length && lines[0].length > 0) {
outputs[outputs.length - 1] = `${outputs[outputs.length - 1]}${lines.shift()!}`;
}
for (const line of lines) {
outputs.push(line);
}
output.items
.filter((opit) => opit.mime === CellOutputMimeTypes.stderr || opit.mime === CellOutputMimeTypes.stdout)
.map((opit) => textDecoder.decode(opit.data))
.forEach(value => {
// Ensure each line is a separate entry in an array (ending with \n).
const lines = value.split('\n');
// If the last item in `outputs` is not empty and the first item in `lines` is not empty, then concate them.
// As they are part of the same line.
if (outputs.length && lines.length && lines[0].length > 0) {
outputs[outputs.length - 1] = `${outputs[outputs.length - 1]}${lines.shift()!}`;
}
for (const line of lines) {
outputs.push(line);
}
});
for (let index = 0; index < (outputs.length - 1); index++) {
outputs[index] = `${outputs[index]}\n`;

View File

@@ -1,63 +0,0 @@
/*---------------------------------------------------------------------------------------------
* Copyright (c) Microsoft Corporation. All rights reserved.
* Licensed under the MIT License. See License.txt in the project root for license information.
*--------------------------------------------------------------------------------------------*/
import type { NotebookCellOutputItem } from 'vscode';
/**
* Given a stream of individual stdout outputs, this function will return the compressed lines, escaping some of the common terminal escape codes.
* E.g. some terminal escape codes would result in the previous line getting cleared, such if we had 3 lines and
* last line contained such a code, then the result string would be just the first two lines.
*/
export function compressOutputItemStreams(mimeType: string, outputs: NotebookCellOutputItem[]) {
// return outputs.find(op => op.mime === mimeType)!.data.buffer;
const buffers: Uint8Array[] = [];
let startAppending = false;
// Pick the first set of outputs with the same mime type.
for (const output of outputs) {
if (output.mime === mimeType) {
if ((buffers.length === 0 || startAppending)) {
buffers.push(output.data);
startAppending = true;
}
} else if (startAppending) {
startAppending = false;
}
}
compressStreamBuffer(buffers);
const totalBytes = buffers.reduce((p, c) => p + c.byteLength, 0);
const combinedBuffer = new Uint8Array(totalBytes);
let offset = 0;
for (const buffer of buffers) {
combinedBuffer.set(buffer, offset);
offset = offset + buffer.byteLength;
}
return combinedBuffer;
}
const MOVE_CURSOR_1_LINE_COMMAND = `${String.fromCharCode(27)}[A`;
const MOVE_CURSOR_1_LINE_COMMAND_BYTES = MOVE_CURSOR_1_LINE_COMMAND.split('').map(c => c.charCodeAt(0));
const LINE_FEED = 10;
function compressStreamBuffer(streams: Uint8Array[]) {
streams.forEach((stream, index) => {
if (index === 0 || stream.length < MOVE_CURSOR_1_LINE_COMMAND.length) {
return;
}
const previousStream = streams[index - 1];
// Remove the previous line if required.
const command = stream.subarray(0, MOVE_CURSOR_1_LINE_COMMAND.length);
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]) {
const lastIndexOfLineFeed = previousStream.lastIndexOf(LINE_FEED);
if (lastIndexOfLineFeed === -1) {
return;
}
streams[index - 1] = previousStream.subarray(0, lastIndexOfLineFeed);
streams[index] = stream.subarray(MOVE_CURSOR_1_LINE_COMMAND.length);
}
});
return streams;
}