Support moving/adding/deleting cells in notebook intellisense (#15208)

rchiodo · web-flow · commit ae6abc0aa6cf · 2021-01-25T15:09:14.000-08:00
* Preliminary idea

* Account for range end position on deletes

* Fix moves

* Test passing

* Simplify movement

* Make sure testing

* Potentially handle multiple deletes or insertions

* Review feedback
diff --git a/src/client/jupyter/languageserver/notebookConcatDocument.ts b/src/client/jupyter/languageserver/notebookConcatDocument.ts
@@ -2,7 +2,21 @@
 // Licensed under the MIT License.
 import * as path from 'path';
 import * as uuid from 'uuid/v4';
-import { Disposable, DocumentSelector, EndOfLine, Position, Range, TextDocument, TextLine, Uri } from 'vscode';
+import {
+    Disposable,
+    DocumentSelector,
+    EndOfLine,
+    Event,
+    EventEmitter,
+    Position,
+    Location,
+    Range,
+    TextDocument,
+    TextDocumentChangeEvent,
+    TextLine,
+    Uri,
+} from 'vscode';
+import { isEqual } from 'lodash';
 import { NotebookConcatTextDocument, NotebookCell, NotebookDocument } from 'vscode-proposed';
 import { IVSCodeNotebook } from '../../common/application/types';
 import { IDisposable } from '../../common/types';
@@ -54,6 +68,10 @@ export class NotebookConcatDocument implements TextDocument, IDisposable {
         return this.notebook.cells.map((c) => c.document.lineCount).reduce((p, c) => p + c);
     }
 
+    public get onCellsChanged(): Event<TextDocumentChangeEvent> {
+        return this.onCellsChangedEmitter.event;
+    }
+
     public firedOpen = false;
 
     public firedClose = false;
@@ -68,6 +86,10 @@ export class NotebookConcatDocument implements TextDocument, IDisposable {
 
     private onDidChangeSubscription: Disposable;
 
+    private cellTracking: { uri: Uri; lineCount: number; length: number }[] = [];
+
+    private onCellsChangedEmitter = new EventEmitter<TextDocumentChangeEvent>();
+
     constructor(public notebook: NotebookDocument, notebookApi: IVSCodeNotebook, selector: DocumentSelector) {
         const dir = path.dirname(notebook.uri.fsPath);
         // Note: Has to be different than the prefix for old notebook editor (HiddenFileFormat) so
@@ -76,6 +98,7 @@ export class NotebookConcatDocument implements TextDocument, IDisposable {
         this.dummyUri = Uri.file(this.dummyFilePath);
         this.concatDocument = notebookApi.createConcatTextDocument(notebook, selector);
         this.onDidChangeSubscription = this.concatDocument.onDidChange(this.onDidChange, this);
+        this.updateCellTracking();
     }
 
     public dispose(): void {
@@ -139,7 +162,133 @@ export class NotebookConcatDocument implements TextDocument, IDisposable {
         return this.notebook.cells.find((c) => c.uri === location.uri);
     }
 
+    private updateCellTracking() {
+        this.cellTracking = [];
+        this.notebook.cells.forEach((c) => {
+            // Compute end position from number of lines in a cell
+            const cellText = c.document.getText();
+            const lines = cellText.splitLines({ trim: false });
+
+            this.cellTracking.push({
+                uri: c.uri,
+                length: cellText.length + 1, // \n is included concat length
+                lineCount: lines.length,
+            });
+        });
+    }
+
     private onDidChange() {
         this._version += 1;
+        const newUris = this.notebook.cells.map((c) => c.uri.toString());
+        const oldUris = this.cellTracking.map((c) => c.uri.toString());
+
+        // See if number of cells or cell positions changed
+        if (this.cellTracking.length < this.notebook.cells.length) {
+            this.raiseCellInsertions(oldUris);
+        } else if (this.cellTracking.length > this.notebook.cells.length) {
+            this.raiseCellDeletions(newUris, oldUris);
+        } else if (!isEqual(oldUris, newUris)) {
+            this.raiseCellMovement();
+        }
+        this.updateCellTracking();
+    }
+
+    private getPositionOfCell(cellUri: Uri): Position {
+        return this.concatDocument.positionAt(new Location(cellUri, new Position(0, 0)));
+    }
+
+    public getEndPosition(): Position {
+        if (this.notebook.cells.length > 0) {
+            const finalCell = this.notebook.cells[this.notebook.cells.length - 1];
+            const start = this.getPositionOfCell(finalCell.uri);
+            const lines = finalCell.document.getText().splitLines({ trim: false });
+            return new Position(start.line + lines.length, 0);
+        }
+        return new Position(0, 0);
+    }
+
+    private raiseCellInsertions(oldUris: string[]) {
+        // One or more cells were added. Add a change event for each
+        const insertions = this.notebook.cells.filter((c) => !oldUris.includes(c.uri.toString()));
+
+        const changes = insertions.map((insertion) => {
+            // Figure out the position of the item. This is where we're inserting the cell
+            // Note: The first insertion will line up with the old cell at this position
+            // The second or other insertions will line up with their new positions.
+            const position = this.getPositionOfCell(insertion.uri);
+
+            // Text should be the contents of the new cell plus the '\n'
+            const text = `${insertion.document.getText()}\n`;
+
+            return {
+                text,
+                range: new Range(position, position),
+                rangeLength: 0,
+                rangeOffset: 0,
+            };
+        });
+
+        // Send all of the changes
+        this.onCellsChangedEmitter.fire({
+            document: this,
+            contentChanges: changes,
+        });
+    }
+
+    private raiseCellDeletions(newUris: string[], oldUris: string[]) {
+        // cells were deleted. Figure out which ones
+        const oldIndexes: number[] = [];
+        oldUris.forEach((o, i) => {
+            if (!newUris.includes(o)) {
+                oldIndexes.push(i);
+            }
+        });
+        const changes = oldIndexes.map((index) => {
+            // Figure out the position of the item in the new list
+            const position =
+                index < newUris.length ? this.getPositionOfCell(this.notebook.cells[index].uri) : this.getEndPosition();
+
+            // Length should be old length
+            const { length } = this.cellTracking[index];
+
+            // Range should go from new position to end of old position
+            const endPosition = new Position(position.line + this.cellTracking[index].lineCount, 0);
+
+            // Turn this cell into a change event.
+            return {
+                text: '',
+                range: new Range(position, endPosition),
+                rangeLength: length,
+                rangeOffset: 0,
+            };
+        });
+
+        // Send the event
+        this.onCellsChangedEmitter.fire({
+            document: this,
+            contentChanges: changes,
+        });
+    }
+
+    private raiseCellMovement() {
+        // When moving, just replace everything. Simpler this way. Might this
+        // cause unknown side effects? Don't think so.
+        this.onCellsChangedEmitter.fire({
+            document: this,
+            contentChanges: [
+                {
+                    text: this.concatDocument.getText(),
+                    range: new Range(
+                        new Position(0, 0),
+                        new Position(
+                            this.cellTracking.reduce((p, c) => p + c.lineCount, 0),
+                            0,
+                        ),
+                    ),
+                    rangeLength: this.cellTracking.reduce((p, c) => p + c.length, 0),
+                    rangeOffset: 0,
+                },
+            ],
+        });
     }
 }
diff --git a/src/client/jupyter/languageserver/notebookConverter.ts b/src/client/jupyter/languageserver/notebookConverter.ts
@@ -15,6 +15,8 @@ import {
     DocumentLink,
     DocumentSelector,
     DocumentSymbol,
+    Event,
+    EventEmitter,
     Hover,
     Location,
     LocationLink,
@@ -44,12 +46,18 @@ function toPosition(positionLike: Position): Position {
 */
 
 export class NotebookConverter implements Disposable {
+    public get onDidChangeCells(): Event<TextDocumentChangeEvent> {
+        return this.onDidChangeCellsEmitter.event;
+    }
+
     private activeDocuments: Map<string, NotebookConcatDocument> = new Map<string, NotebookConcatDocument>();
 
     private activeDocumentsOutgoingMap: Map<string, NotebookConcatDocument> = new Map<string, NotebookConcatDocument>();
 
     private disposables: Disposable[] = [];
 
+    private onDidChangeCellsEmitter = new EventEmitter<TextDocumentChangeEvent>();
+
     constructor(
         private api: IVSCodeNotebook,
         private fs: IFileSystem,
@@ -627,6 +635,7 @@ export class NotebookConverter implements Disposable {
                 throw new Error(`Invalid uri, not a notebook: ${uri.fsPath}`);
             }
             result = new NotebookConcatDocument(doc, this.api, this.cellSelector);
+            result.onCellsChanged((e) => this.onDidChangeCellsEmitter.fire(e), undefined, this.disposables);
             this.activeDocuments.set(key, result);
             this.activeDocumentsOutgoingMap.set(NotebookConverter.getDocumentKey(result.uri), result);
         }
diff --git a/src/client/jupyter/languageserver/notebookMiddlewareAddon.ts b/src/client/jupyter/languageserver/notebookMiddlewareAddon.ts
@@ -76,6 +76,8 @@ import { NotebookConverter } from './notebookConverter';
 export class NotebookMiddlewareAddon implements Middleware, Disposable {
     private converter: NotebookConverter;
 
+    private didChangeCellsDisposable: Disposable;
+
     constructor(
         notebookApi: IVSCodeNotebook,
         private readonly getClient: () => LanguageClient | undefined,
@@ -84,9 +86,11 @@ export class NotebookMiddlewareAddon implements Middleware, Disposable {
         notebookFileRegex: RegExp,
     ) {
         this.converter = new NotebookConverter(notebookApi, fs, cellSelector, notebookFileRegex);
+        this.didChangeCellsDisposable = this.converter.onDidChangeCells(this.onDidChangeCells.bind(this));
     }
 
     public dispose(): void {
+        this.didChangeCellsDisposable.dispose();
         this.converter.dispose();
     }
 
@@ -481,4 +485,16 @@ export class NotebookMiddlewareAddon implements Middleware, Disposable {
         const newDiagMapping = this.converter.toIncomingDiagnosticsMap(uri, diagnostics);
         [...newDiagMapping.keys()].forEach((k) => next(k, newDiagMapping.get(k)!));
     }
+
+    private onDidChangeCells(e: TextDocumentChangeEvent) {
+        // This event fires when the user moves, deletes, or inserts cells into the concatenated document
+        // Since this doesn't fire a change event (since a document wasn't changed), we have to make one ourselves.
+
+        // Note: The event should already be setup to be an outgoing event. It's from the point of view of the concatenated document.
+        const client = this.getClient();
+        if (client) {
+            const params = client.code2ProtocolConverter.asChangeTextDocumentParams(e);
+            client.sendNotification(DidChangeTextDocumentNotification.type, params);
+        }
+    }
 }
diff --git a/src/test/initialize.ts b/src/test/initialize.ts
@@ -86,7 +86,7 @@ async function closeWindowsInteral() {
         // Lets not waste too much time.
         const timer = setTimeout(() => {
             reject(new Error("Command 'workbench.action.closeAllEditors' timed out"));
-        }, 15000);
+        }, 2000);
         vscode.commands.executeCommand('workbench.action.closeAllEditors').then(
             () => {
                 clearTimeout(timer);
diff --git a/src/test/insiders/languageServer.insiders.test.ts b/src/test/insiders/languageServer.insiders.test.ts
@@ -7,6 +7,7 @@ import * as assert from 'assert';
 import { expect } from 'chai';
 import * as path from 'path';
 import * as vscode from 'vscode';
+import { PYTHON_LANGUAGE } from '../../client/common/constants';
 import { updateSetting } from '../common';
 import { EXTENSION_ROOT_DIR_FOR_TESTS } from '../constants';
 import { sleep } from '../core';
@@ -94,15 +95,75 @@ suite('Insiders Test: Language Server', () => {
             );
             if (locations && locations.length > 0) {
                 expect(locations![0].uri.fsPath).to.contain(path.basename(notebookDefinitions));
+
+                // Insert a new cell
+                const activeEditor = vscode.notebook.activeNotebookEditor;
+                expect(activeEditor).not.to.be.equal(undefined, 'Active editor not found in notebook');
+                await activeEditor!.edit((edit) => {
+                    edit.replaceCells(0, 0, [
+                        {
+                            cellKind: vscode.CellKind.Code,
+                            language: PYTHON_LANGUAGE,
+                            source: 'x = 4',
+                            metadata: {
+                                hasExecutionOrder: false,
+                            },
+                            outputs: [],
+                        },
+                    ]);
+                });
+
+                // Wait a bit to get diagnostics
+                await sleep(1_000);
+
+                // Make sure no error diagnostics
+                let diagnostics = vscode.languages.getDiagnostics(activeEditor!.document.uri);
+                expect(diagnostics).to.have.lengthOf(0, 'Diagnostics found when shouldnt be');
+
+                // Move the cell
+                await activeEditor!.edit((edit) => {
+                    edit.replaceCells(0, 1, []);
+                    edit.replaceCells(1, 0, [
+                        {
+                            cellKind: vscode.CellKind.Code,
+                            language: PYTHON_LANGUAGE,
+                            source: 'x = 4',
+                            metadata: {
+                                hasExecutionOrder: false,
+                            },
+                            outputs: [],
+                        },
+                    ]);
+                });
+
+                // Wait a bit to get diagnostics
+                await sleep(1_000);
+
+                // Make sure no error diagnostics
+                diagnostics = vscode.languages.getDiagnostics(activeEditor!.document.uri);
+                expect(diagnostics).to.have.lengthOf(0, 'Diagnostics found when shouldnt be after move');
+
+                // Delete the cell
+                await activeEditor!.edit((edit) => {
+                    edit.replaceCells(1, 1, []);
+                });
+
+                // Wait a bit to get diagnostics
+                await sleep(1_000);
+
+                // Make sure no error diagnostics
+                diagnostics = vscode.languages.getDiagnostics(activeEditor!.document.uri);
+                expect(diagnostics).to.have.lengthOf(0, 'Diagnostics found when shouldnt be after delete');
                 tested = true;
+
                 break;
             } else {
                 // Wait for LS to start.
                 await sleep(5_000);
             }
         }
         if (!tested) {
-            assert.fail('Failled to test definitions');
+            assert.fail('Failled to test notebooks');
         }
     });
 });
