Skip to content

Commit 79c4954

Browse files
amungeramunger
committed
track output buffer lengths per version
1 parent fbd39c1 commit 79c4954

File tree

3 files changed

+152
-9
lines changed

3 files changed

+152
-9
lines changed

‎src/vs/workbench/contrib/notebook/common/model/notebookCellOutputTextModel.ts‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
* Licensed under the MIT License. See License.txt in the project root for license information.
44
*--------------------------------------------------------------------------------------------*/
55

6+
import { VSBuffer } from 'vs/base/common/buffer';
67
import { Emitter } from 'vs/base/common/event';
78
import { Disposable } from 'vs/base/common/lifecycle';
89
import { ICellOutput, IOutputDto, IOutputItemDto, compressOutputItemStreams, isTextStreamMime } from 'vs/workbench/contrib/notebook/common/notebookCommon';
@@ -37,19 +38,45 @@ export class NotebookCellOutputTextModel extends Disposable implements ICellOutp
3738
}
3839

3940
replaceData(rawData: IOutputDto) {
41+
this.versionedBufferLengths = {};
4042
this._rawOutput = rawData;
4143
this.optimizeOutputItems();
4244
this._versionId = this._versionId + 1;
4345
this._onDidChangeData.fire();
4446
}
4547

4648
appendData(items: IOutputItemDto[]) {
49+
this.trackBufferLengths();
4750
this._rawOutput.outputs.push(...items);
4851
this.optimizeOutputItems();
4952
this._versionId = this._versionId + 1;
5053
this._onDidChangeData.fire();
5154
}
5255

56+
private trackBufferLengths() {
57+
this.outputs.forEach(output => {
58+
if (isTextStreamMime(output.mime)) {
59+
if (!this.versionedBufferLengths[output.mime]) {
60+
this.versionedBufferLengths[output.mime] = {};
61+
}
62+
this.versionedBufferLengths[output.mime][this.versionId] = output.data.byteLength;
63+
}
64+
});
65+
}
66+
67+
// mime: versionId: buffer length
68+
private versionedBufferLengths: Record<string, Record<number, number>> = {};
69+
70+
appendedSinceVersion(versionId: number, mime: string): VSBuffer | undefined {
71+
const bufferLength = this.versionedBufferLengths[mime]?.[versionId];
72+
const output = this.outputs.find(output => output.mime === mime);
73+
if (bufferLength && output) {
74+
return output.data.slice(bufferLength);
75+
}
76+
77+
return undefined;
78+
}
79+
5380
private optimizeOutputItems() {
5481
if (this.outputs.length > 1 && this.outputs.every(item => isTextStreamMime(item.mime))) {
5582
// Look for the mimes in the items, and keep track of their order.

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,7 @@ export interface ICellOutput {
217217
onDidChangeData: Event<void>;
218218
replaceData(items: IOutputDto): void;
219219
appendData(items: IOutputItemDto[]): void;
220+
appendedSinceVersion(versionId: number, mime: string): VSBuffer | undefined;
220221
}
221222

222223
export interface CellInternalMetadataChangedEvent {

‎src/vs/workbench/contrib/notebook/test/browser/notebookTextModel.test.ts‎

Lines changed: 124 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -311,6 +311,9 @@ suite('NotebookTextModel', () => {
311311
);
312312
});
313313

314+
const stdOutMime = 'application/vnd.code.notebook.stdout';
315+
const stdErrMime = 'application/vnd.code.notebook.stderr';
316+
314317
test('appending streaming outputs', async function () {
315318
await withTestNotebook(
316319
[
@@ -326,20 +329,21 @@ suite('NotebookTextModel', () => {
326329
append: true,
327330
outputs: [{
328331
outputId: 'append1',
329-
outputs: [{ mime: 'application/vnd.code.notebook.stdout', data: valueBytesFromString('append 1') }]
332+
outputs: [{ mime: stdOutMime, data: valueBytesFromString('append 1') }]
330333
}]
331334
}], true, undefined, () => undefined, undefined, true);
332335
const [output] = textModel.cells[0].outputs;
333-
assert.strictEqual(output.versionId, 0, 'initial output version is 0');
336+
assert.strictEqual(output.versionId, 0, 'initial output version should be 0');
334337

335338
textModel.applyEdits([
336339
{
337340
editType: CellEditType.OutputItems,
338341
append: true,
339342
outputId: 'append1',
340-
items: [{
341-
mime: 'application/vnd.code.notebook.stdout', data: valueBytesFromString('append 2')
342-
}]
343+
items: [
344+
{ mime: stdOutMime, data: valueBytesFromString('append 2') },
345+
{ mime: stdOutMime, data: valueBytesFromString('append 3') }
346+
]
343347
}], true, undefined, () => undefined, undefined, true);
344348
assert.strictEqual(output.versionId, 1, 'version should bump per append');
345349

@@ -348,17 +352,128 @@ suite('NotebookTextModel', () => {
348352
editType: CellEditType.OutputItems,
349353
append: true,
350354
outputId: 'append1',
351-
items: [{
352-
mime: 'application/vnd.code.notebook.stdout', data: valueBytesFromString('append 3')
353-
}]
355+
items: [
356+
{ mime: stdOutMime, data: valueBytesFromString('append 4') },
357+
{ mime: stdOutMime, data: valueBytesFromString('append 5') }
358+
]
354359
}], true, undefined, () => undefined, undefined, true);
355360
assert.strictEqual(output.versionId, 2, 'version should bump per append');
356361

357362
assert.strictEqual(textModel.cells.length, 1);
358363
assert.strictEqual(textModel.cells[0].outputs.length, 1, 'has 1 output');
359364
assert.strictEqual(output.outputId, 'append1');
360365
assert.strictEqual(output.outputs.length, 1, 'outputs are compressed');
361-
assert.strictEqual(output.outputs[0].data.toString(), 'append 1append 2append 3');
366+
assert.strictEqual(output.outputs[0].data.toString(), 'append 1append 2append 3append 4append 5');
367+
assert.strictEqual(output.appendedSinceVersion(0, stdOutMime)?.toString(), 'append 2append 3append 4append 5');
368+
assert.strictEqual(output.appendedSinceVersion(1, stdOutMime)?.toString(), 'append 4append 5');
369+
assert.strictEqual(output.appendedSinceVersion(2, stdOutMime), undefined);
370+
assert.strictEqual(output.appendedSinceVersion(2, stdErrMime), undefined);
371+
}
372+
);
373+
});
374+
375+
test('replacing streaming outputs', async function () {
376+
await withTestNotebook(
377+
[
378+
['var a = 1;', 'javascript', CellKind.Code, [], {}],
379+
],
380+
(editor) => {
381+
const textModel = editor.textModel;
382+
383+
textModel.applyEdits([
384+
{
385+
index: 0,
386+
editType: CellEditType.Output,
387+
append: true,
388+
outputs: [{
389+
outputId: 'append1',
390+
outputs: [{ mime: stdOutMime, data: valueBytesFromString('append 1') }]
391+
}]
392+
}], true, undefined, () => undefined, undefined, true);
393+
const [output] = textModel.cells[0].outputs;
394+
assert.strictEqual(output.versionId, 0, 'initial output version should be 0');
395+
396+
textModel.applyEdits([
397+
{
398+
editType: CellEditType.OutputItems,
399+
append: true,
400+
outputId: 'append1',
401+
items: [{
402+
mime: stdOutMime, data: valueBytesFromString('append 2')
403+
}]
404+
}], true, undefined, () => undefined, undefined, true);
405+
assert.strictEqual(output.versionId, 1, 'version should bump per append');
406+
407+
textModel.applyEdits([
408+
{
409+
editType: CellEditType.OutputItems,
410+
append: false,
411+
outputId: 'append1',
412+
items: [{
413+
mime: stdOutMime, data: valueBytesFromString('replace 3')
414+
}]
415+
}], true, undefined, () => undefined, undefined, true);
416+
assert.strictEqual(output.versionId, 2, 'version should bump per replace');
417+
418+
textModel.applyEdits([
419+
{
420+
editType: CellEditType.OutputItems,
421+
append: true,
422+
outputId: 'append1',
423+
items: [{
424+
mime: stdOutMime, data: valueBytesFromString('append 4')
425+
}]
426+
}], true, undefined, () => undefined, undefined, true);
427+
assert.strictEqual(output.versionId, 3, 'version should bump per append');
428+
429+
assert.strictEqual(output.outputs[0].data.toString(), 'replace 3append 4');
430+
assert.strictEqual(output.appendedSinceVersion(0, stdOutMime), undefined,
431+
'replacing output should clear out previous versioned output buffers');
432+
assert.strictEqual(output.appendedSinceVersion(1, stdOutMime), undefined,
433+
'replacing output should clear out previous versioned output buffers');
434+
assert.strictEqual(output.appendedSinceVersion(2, stdOutMime)?.toString(), 'append 4');
435+
}
436+
);
437+
});
438+
439+
test('appending multiple different mime streaming outputs', async function () {
440+
await withTestNotebook(
441+
[
442+
['var a = 1;', 'javascript', CellKind.Code, [], {}],
443+
],
444+
(editor) => {
445+
const textModel = editor.textModel;
446+
447+
textModel.applyEdits([
448+
{
449+
index: 0,
450+
editType: CellEditType.Output,
451+
append: true,
452+
outputs: [{
453+
outputId: 'append1',
454+
outputs: [
455+
{ mime: stdOutMime, data: valueBytesFromString('stdout 1') },
456+
{ mime: stdErrMime, data: valueBytesFromString('stderr 1') }
457+
]
458+
}]
459+
}], true, undefined, () => undefined, undefined, true);
460+
const [output] = textModel.cells[0].outputs;
461+
assert.strictEqual(output.versionId, 0, 'initial output version should be 0');
462+
463+
textModel.applyEdits([
464+
{
465+
editType: CellEditType.OutputItems,
466+
append: true,
467+
outputId: 'append1',
468+
items: [
469+
{ mime: stdOutMime, data: valueBytesFromString('stdout 2') },
470+
{ mime: stdErrMime, data: valueBytesFromString('stderr 2') }
471+
]
472+
}], true, undefined, () => undefined, undefined, true);
473+
assert.strictEqual(output.versionId, 1, 'version should bump per replace');
474+
475+
assert.strictEqual(output.appendedSinceVersion(0, stdErrMime)?.toString(), 'stderr 2');
476+
assert.strictEqual(output.appendedSinceVersion(0, stdOutMime)?.toString(), 'stdout 2');
362477
}
363478
);
364479
});

0 commit comments

Comments
 (0)