Update VSC Notebook API (#12333)

Author: DonJayamanne

src/client/common/application/notebook.ts

@@ -34,6 +34,9 @@ export class VSCodeNotebook implements IVSCodeNotebook {
     public get onDidCloseNotebookDocument(): Event<NotebookDocument> {
         return this.notebook.onDidCloseNotebookDocument;
     }
+    public get notebookDocuments(): ReadonlyArray<NotebookDocument> {
+        return this.notebook.notebookDocuments;
+    }
     public get notebookEditors() {
         return this.notebook.visibleNotebookEditors;
     }

src/client/common/application/types.ts

@@ -1522,6 +1522,7 @@ export type NotebookCellChangedEvent =
     | NotebookCellLanguageChangeEvent;
 export const IVSCodeNotebook = Symbol('IVSCodeNotebook');
 export interface IVSCodeNotebook {
+    readonly notebookDocuments: ReadonlyArray<NotebookDocument>;
     readonly onDidOpenNotebookDocument: Event<NotebookDocument>;
     readonly onDidCloseNotebookDocument: Event<NotebookDocument>;
     readonly onDidChangeActiveNotebookEditor: Event<NotebookEditor | undefined>;

src/client/datascience/notebook/cellEditSyncService.ts

@@ -109,7 +109,7 @@ export class CellEditSyncService implements IExtensionSingleActivationService, I
             return;
         }
 
-        this.mappedDocuments.set(cellDocument, { model: editor.model, cellId: getOriginalCellId(document, cell)! });
+        this.mappedDocuments.set(cellDocument, { model: editor.model, cellId: getOriginalCellId(cell)! });
         return this.mappedDocuments.get(cellDocument);
     }
 }

src/client/datascience/notebook/executionService.ts

@@ -179,7 +179,7 @@ export class NotebookExecutionService implements INotebookExecutionService {
                         .flatMap((item) => (item.data.outputs as unknown) as nbformat.IOutput[])
                         .filter((output) => !hasTransientOutputForAnotherCell(output));
 
-                    const notebookCellModel = findMappedNotebookCellModel(document, cell, model.cells);
+                    const notebookCellModel = findMappedNotebookCellModel(cell, model.cells);
 
                     // Set execution count, all messages should have it
                     if (
@@ -206,7 +206,7 @@ export class NotebookExecutionService implements INotebookExecutionService {
                     cell.metadata.statusMessage = '';
 
                     // Update metadata in our model.
-                    const notebookCellModel = findMappedNotebookCellModel(document, cell, model.cells);
+                    const notebookCellModel = findMappedNotebookCellModel(cell, model.cells);
                     updateCellExecutionTimes(
                         notebookCellModel,
                         cell.metadata.runStartTime,

src/client/datascience/notebook/helpers/cellMappers.ts

@@ -25,16 +25,16 @@ type ICellId = string;
 const cellMapsByNotebookDocument = new WeakMap<NotebookDocument, Map<VSCodeCellUri, ICellId>>();
 const mapOfCellsToNotebookCells = new WeakMap<ICell, VSCodeCellUri>();
 
-export function mapVSCNotebookCellToCellModel(document: NotebookDocument, vscCell: NotebookCell, cell: ICell) {
-    if (!cellMapsByNotebookDocument.has(document)) {
-        cellMapsByNotebookDocument.set(document, new Map<VSCodeCellUri, ICellId>());
+export function mapVSCNotebookCellToCellModel(vscCell: NotebookCell, cell: ICell) {
+    if (!cellMapsByNotebookDocument.has(vscCell.notebook)) {
+        cellMapsByNotebookDocument.set(vscCell.notebook, new Map<VSCodeCellUri, ICellId>());
     }
-    cellMapsByNotebookDocument.get(document)!.set(vscCell.uri.toString(), cell.id);
+    cellMapsByNotebookDocument.get(vscCell.notebook)!.set(vscCell.uri.toString(), cell.id);
     mapOfCellsToNotebookCells.set(cell, vscCell.uri.toString());
 }
 
-export function getOriginalCellId(document: NotebookDocument, cell: NotebookCell): string | undefined {
-    const map = cellMapsByNotebookDocument.get(document);
+export function getOriginalCellId(cell: NotebookCell): string | undefined {
+    const map = cellMapsByNotebookDocument.get(cell.notebook);
     if (!map) {
         return;
     }
@@ -61,7 +61,7 @@ export function mapVSCNotebookCellsToNotebookCellModels(document: NotebookDocume
 
     document.cells.forEach((vscCell, index) => {
         const cell = model.cells[index];
-        mapVSCNotebookCellToCellModel(document, vscCell, cell);
+        mapVSCNotebookCellToCellModel(vscCell, cell);
     });
 }
 
@@ -74,10 +74,10 @@ export function findMappedNotebookCell(source: ICell, cells: NotebookCell[]): No
     return found[0];
 }
 
-export function findMappedNotebookCellModel(document: NotebookDocument, source: NotebookCell, cells: ICell[]): ICell {
+export function findMappedNotebookCellModel(source: NotebookCell, cells: ICell[]): ICell {
     // If so, then we have a problem.
-    const found = cells.filter((cell) => cell.id === getOriginalCellId(document, source));
-    assert.ok(found.length, `ICell not found, for CellId = ${getOriginalCellId(document, source)} in ${source}`);
+    const found = cells.filter((cell) => cell.id === getOriginalCellId(source));
+    assert.ok(found.length, `ICell not found, for CellId = ${getOriginalCellId(source)} in ${source}`);
 
     return found[0];
 }

src/client/datascience/notebook/helpers/cellUpdateHelpers.ts

@@ -59,7 +59,7 @@ function clearCellOutput(change: NotebookCellOutputsChangeEvent, model: INoteboo
 
     // If a cell has been cleared, then clear the corresponding ICell (cell in INotebookModel).
     change.cells.forEach((vscCell) => {
-        const cell = findMappedNotebookCellModel(change.document, vscCell, model.cells);
+        const cell = findMappedNotebookCellModel(vscCell, model.cells);
         cell.data.outputs = [];
         updateVSCNotebookCellMetadata(vscCell.metadata, cell);
         model.update({
@@ -73,7 +73,7 @@ function clearCellOutput(change: NotebookCellOutputsChangeEvent, model: INoteboo
 }
 
 function changeCellLanguage(change: NotebookCellLanguageChangeEvent, model: INotebookModel) {
-    const cellModel = findMappedNotebookCellModel(change.document, change.cell, model.cells);
+    const cellModel = findMappedNotebookCellModel(change.cell, model.cells);
 
     // VSC fires event if changing cell language from markdown to markdown.
     // https://github.com/microsoft/vscode/issues/98836
@@ -93,7 +93,7 @@ function changeCellLanguage(change: NotebookCellLanguageChangeEvent, model: INot
 
     // Create a new cell & replace old one.
     const oldCellIndex = model.cells.indexOf(cellModel);
-    model.cells[oldCellIndex] = createCellFromVSCNotebookCell(change.document, change.cell, model);
+    model.cells[oldCellIndex] = createCellFromVSCNotebookCell(change.cell, model);
 }
 
 function handleChangesToCells(change: NotebookCellsChangeEvent, model: INotebookModel) {
@@ -149,7 +149,7 @@ function isCellInsertion(change: NotebookCellsChangeEvent) {
 function handleCellMove(change: NotebookCellsChangeEvent, model: INotebookModel) {
     assert.equal(change.changes.length, 2, 'When moving cells we must have only 2 changes');
     const [, insertChange] = change.changes;
-    const cellToSwap = findMappedNotebookCellModel(change.document, insertChange.items[0]!, model.cells);
+    const cellToSwap = findMappedNotebookCellModel(insertChange.items[0]!, model.cells);
     const cellToSwapWith = model.cells[insertChange.start];
     assert.notEqual(cellToSwap, cellToSwapWith, 'Cannot swap cell with the same cell');
 
@@ -162,7 +162,7 @@ function handleCellInsertion(change: NotebookCellsChangeEvent, model: INotebookM
     assert.equal(change.changes[0].items.length, 1, 'Insertion of more than 1 cell is not supported');
     const insertChange = change.changes[0];
     const cell = change.changes[0].items[0];
-    const newCell = createCellFromVSCNotebookCell(change.document, cell, model);
+    const newCell = createCellFromVSCNotebookCell(cell, model);
     model.cells.splice(insertChange.start, 0, newCell);
 }
 function handleCellDelete(change: NotebookCellsChangeEvent, model: INotebookModel) {

src/client/datascience/notebook/helpers/helpers.ts

@@ -14,8 +14,7 @@ import type {
     NotebookCell,
     NotebookCellData,
     NotebookCellMetadata,
-    NotebookData,
-    NotebookDocument
+    NotebookData
 } from 'vscode-proposed';
 import { NotebookCellRunState } from '../../../../../typings/vscode-proposed';
 import {
@@ -76,11 +75,7 @@ export function notebookModelToVSCNotebookData(model: INotebookModel): NotebookD
         }
     };
 }
-export function createCellFromVSCNotebookCell(
-    document: NotebookDocument,
-    vscCell: NotebookCell,
-    model: INotebookModel
-): ICell {
+export function createCellFromVSCNotebookCell(vscCell: NotebookCell, model: INotebookModel): ICell {
     const cell = (() => {
         if (vscCell.cellKind === vscodeNotebookEnums.CellKind.Markdown) {
             return {
@@ -112,7 +107,7 @@ export function createCellFromVSCNotebookCell(
         };
     }
     // Ensure we add the cell id of the new cell to the VSC cell to map into ours.
-    mapVSCNotebookCellToCellModel(document, vscCell, cell);
+    mapVSCNotebookCellToCellModel(vscCell, cell);
 
     return cell;
 }

src/test/datascience/notebook/contentProvider.ds.test.ts

@@ -68,11 +68,7 @@ suite('DataScience - VSCode Notebook - (Open)', function () {
             const vscCell = notebook.cells[index];
             const cell = model.cells[index];
             // Given a VS Code Cell, we should be able to find the corresponding ICell.
-            assert.equal(
-                cell,
-                findMappedNotebookCellModel(notebook, vscCell, model.cells),
-                'Could not find mapped ICell'
-            );
+            assert.equal(cell, findMappedNotebookCellModel(vscCell, model.cells), 'Could not find mapped ICell');
 
             // Given an ICell, we should be able to find the corresponding VS Code Cell.
             assert.equal(vscCell, findMappedNotebookCell(cell, notebook.cells), 'Could not find mapped NotebookCell');

src/test/datascience/notebook/helper.ts

@@ -77,7 +77,7 @@ export async function insertPythonCell(source: string, index: number = 0) {
             await waitForCondition(
                 async () =>
                     vscEditor!.document.cells.every((cell) =>
-                        findMappedNotebookCellModel(vscEditor!.document, cell, nbEditor!.model!.cells)
+                        findMappedNotebookCellModel(cell, nbEditor!.model!.cells)
                     ),
                 1_000,
                 'Cell not assigned a cell Id'

types/vscode-proposed/index.d.ts

@@ -45,6 +45,13 @@ export interface CellErrorOutput {
     traceback: string[];
 }
 
+export interface NotebookCellOutputMetadata {
+    /**
+     * Additional attributes of a cell metadata.
+     */
+    custom?: { [key: string]: any };
+}
+
 export interface CellDisplayOutput {
     outputKind: CellOutputKind.Rich;
     /**
@@ -65,6 +72,8 @@ export interface CellDisplayOutput {
      * }
      */
     data: { [key: string]: any };
+
+    readonly metadata?: NotebookCellOutputMetadata;
 }
 
 export type CellOutput = CellStreamOutput | CellErrorOutput | CellDisplayOutput;
@@ -132,11 +141,10 @@ export interface NotebookCellMetadata {
 }
 
 export interface NotebookCell {
+    readonly notebook: NotebookDocument;
     readonly uri: Uri;
     readonly cellKind: CellKind;
     readonly document: TextDocument;
-    // API remove `source` or doc it as shorthand for document.getText()
-    readonly source: string;
     language: string;
     outputs: CellOutput[];
     metadata: NotebookCellMetadata;
@@ -184,6 +192,7 @@ export interface NotebookDocumentMetadata {
 export interface NotebookDocument {
     readonly uri: Uri;
     readonly fileName: string;
+    readonly viewType: string;
     readonly isDirty: boolean;
     readonly cells: NotebookCell[];
     languages: string[];
@@ -286,6 +295,7 @@ export interface NotebookOutputRenderer {
 export interface NotebookCellsChangeData {
     readonly start: number;
     readonly deletedCount: number;
+    readonly deletedItems: NotebookCell[];
     readonly items: NotebookCell[];
 }
 
@@ -379,12 +389,15 @@ export namespace notebook {
 
     export const onDidOpenNotebookDocument: Event<NotebookDocument>;
     export const onDidCloseNotebookDocument: Event<NotebookDocument>;
+
+    /**
+     * All currently known notebook documents.
+     */
+    export const notebookDocuments: ReadonlyArray<NotebookDocument>;
+
     export let visibleNotebookEditors: NotebookEditor[];
     export const onDidChangeVisibleNotebookEditors: Event<NotebookEditor[]>;
 
-    // remove activeNotebookDocument, now that there is activeNotebookEditor.document
-    export let activeNotebookDocument: NotebookDocument | undefined;
-
     export let activeNotebookEditor: NotebookEditor | undefined;
     export const onDidChangeActiveNotebookEditor: Event<NotebookEditor | undefined>;
     export const onDidChangeNotebookCells: Event<NotebookCellsChangeEvent>;

types/vscode.proposed.d.ts

@@ -46,6 +46,13 @@ declare module 'vscode' {
         traceback: string[];
     }
 
+    export interface NotebookCellOutputMetadata {
+        /**
+         * Additional attributes of a cell metadata.
+         */
+        custom?: { [key: string]: any };
+    }
+
     export interface CellDisplayOutput {
         outputKind: CellOutputKind.Rich;
         /**
@@ -66,6 +73,8 @@ declare module 'vscode' {
          * }
          */
         data: { [key: string]: any };
+
+        readonly metadata?: NotebookCellOutputMetadata;
     }
 
     export type CellOutput = CellStreamOutput | CellErrorOutput | CellDisplayOutput;
@@ -133,11 +142,10 @@ declare module 'vscode' {
     }
 
     export interface NotebookCell {
+        readonly notebook: NotebookDocument;
         readonly uri: Uri;
         readonly cellKind: CellKind;
         readonly document: TextDocument;
-        // API remove `source` or doc it as shorthand for document.getText()
-        readonly source: string;
         language: string;
         outputs: CellOutput[];
         metadata: NotebookCellMetadata;
@@ -185,6 +193,7 @@ declare module 'vscode' {
     export interface NotebookDocument {
         readonly uri: Uri;
         readonly fileName: string;
+        readonly viewType: string;
         readonly isDirty: boolean;
         readonly cells: NotebookCell[];
         languages: string[];
@@ -287,6 +296,7 @@ declare module 'vscode' {
     export interface NotebookCellsChangeData {
         readonly start: number;
         readonly deletedCount: number;
+        readonly deletedItems: NotebookCell[];
         readonly items: NotebookCell[];
     }
 
@@ -384,12 +394,15 @@ declare module 'vscode' {
 
         export const onDidOpenNotebookDocument: Event<NotebookDocument>;
         export const onDidCloseNotebookDocument: Event<NotebookDocument>;
+
+        /**
+         * All currently known notebook documents.
+         */
+        export const notebookDocuments: ReadonlyArray<NotebookDocument>;
+
         export let visibleNotebookEditors: NotebookEditor[];
         export const onDidChangeVisibleNotebookEditors: Event<NotebookEditor[]>;
 
-        // remove activeNotebookDocument, now that there is activeNotebookEditor.document
-        export let activeNotebookDocument: NotebookDocument | undefined;
-
         export let activeNotebookEditor: NotebookEditor | undefined;
         export const onDidChangeActiveNotebookEditor: Event<NotebookEditor | undefined>;
         export const onDidChangeNotebookCells: Event<NotebookCellsChangeEvent>;