Handle moving cells with custom message

.vscode/launch.json

@@ -11,7 +11,7 @@
             "runtimeExecutable": "${execPath}",
             "args": [
                 "${workspaceRoot}/src/test",
-                "--enable-proposed-api=vscode-jupyter-lsp-middleware", // Name of test extension
+                "--enable-proposed-api=ms-vscode.vscode-jupyter-lsp-middleware", // Name of test extension
                 "--extensionDevelopmentPath=${workspaceFolder:vscode-jupyter-lsp-middleware}/out/test", // Location where extension.ts goes
                 "--extensionTestsPath=${workspaceFolder:vscode-jupyter-lsp-middleware}/out/test/suite/index", // Location of test files,
                 "--skip-welcome",

scripts/prepareTest.js

@@ -1,5 +1,5 @@
 const fs = require('fs');
-
+console.log('Copying package.json to out folder')
 if (!fs.existsSync('./out')) {
   fs.mkdirSync('./out');
   fs.mkdirSync('./out/test');

src/common/types.ts

@@ -1,25 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import {
-    DocumentSelector,
-    NotebookDocument,
-    Event,
-    NotebookConcatTextDocument,
-    Disposable,
-    Position,
-    TextDocument,
-    TextLine,
-    Uri
-} from 'vscode';
-
-export const IVSCodeNotebook = Symbol('IVSCodeNotebook');
-export interface IVSCodeNotebook {
-    readonly notebookDocuments: ReadonlyArray<NotebookDocument>;
-    readonly onDidOpenNotebookDocument: Event<NotebookDocument>;
-    readonly onDidCloseNotebookDocument: Event<NotebookDocument>;
-    createConcatTextDocument(notebook: NotebookDocument, selector?: DocumentSelector): NotebookConcatTextDocument;
-}
+import { Disposable } from 'vscode';
+import * as protocol from 'vscode-languageclient';
 
 export interface IDisposable {
     dispose(): void | undefined;
@@ -27,20 +10,7 @@ export interface IDisposable {
 
 export type TemporaryFile = { filePath: string } & Disposable;
 
-export interface IConcatTextDocument {
-    onDidChange: Event<void>;
-    isClosed: boolean;
-    lineCount: number;
-    languageId: string;
-    isComposeDocumentsAllClosed: boolean;
-    getText(range?: Range): string;
-    contains(uri: Uri): boolean;
-    offsetAt(position: Position): number;
-    positionAt(locationOrOffset: Location | number): Position;
-    validateRange(range: Range): Range;
-    validatePosition(position: Position): Position;
-    locationAt(positionOrRange: Position | Range): Location;
-    lineAt(posOrNumber: Position | number): TextLine;
-    getWordRangeAtPosition(position: Position, regexp?: RegExp | undefined): Range | undefined;
-    getComposeDocuments(): TextDocument[];
-}
+// Type for refresh notebook to pass through LSP
+export type RefreshNotebookEvent = {
+    cells: protocol.DidOpenTextDocumentParams[];
+};

src/notebookConcatDocument.ts

@@ -14,6 +14,7 @@ import {
     regExpLeadsToEndlessLoop
 } from './common/wordHelper';
 import { NotebookConcatLine } from './notebookConcatLine';
+import { RefreshNotebookEvent } from './common/types';
 
 interface ICellRange {
     uri: vscode.Uri;
@@ -199,6 +200,46 @@ export class NotebookConcatDocument implements vscode.TextDocument, vscode.Dispo
         return this.toDidChangeTextDocumentParams(changes);
     }
 
+    public handleRefresh(e: RefreshNotebookEvent): protocol.DidChangeTextDocumentParams | undefined {
+        // Delete all cells and start over. This should only happen for non interactive (you can't move interactive cells at the moment)
+        if (!this._interactiveWindow) {
+            // Track our old full range
+            const from = new vscode.Position(0, 0);
+            const to = this.positionAt(this._contents.length);
+            const oldLength = this._contents.length;
+
+            this._cellRanges = [];
+            this._contents = `${e.cells.map((c) => c.textDocument.text).join('\n')}\n`;
+            this._lines = this.createLines();
+            let startOffset = 0;
+            e.cells.forEach((c) => {
+                const cellUri = vscode.Uri.parse(c.textDocument.uri);
+                this._cellRanges.push({
+                    uri: cellUri,
+                    startOffset,
+                    startLine: this._lines.find((l) => l.offset === startOffset)?.lineNumber || 0,
+                    fragment:
+                        cellUri.scheme === InteractiveInputScheme ? -1 : parseInt(cellUri.fragment.substring(2) || '0'),
+                    endOffset: c.textDocument.text.length + 1 // Account for \n between cells
+                });
+                startOffset = this._cellRanges[this._cellRanges.length - 1].endOffset;
+            });
+
+            // Create one big change
+            const changes: protocol.TextDocumentContentChangeEvent[] = [
+                {
+                    range: this.createSerializableRange(from, to),
+                    rangeOffset: 0,
+                    rangeLength: oldLength,
+                    text: this._contents
+                } as any
+            ];
+
+            return this.toDidChangeTextDocumentParams(changes);
+        }
+        return undefined;
+    }
+
     public dispose() {
         // Do nothing for now.
     }

src/notebookConverter.ts

@@ -108,6 +108,13 @@ export class NotebookConverter implements Disposable {
         return results;
     }
 
+    public handleRefresh(notebook: NotebookDocument) {
+        // Find the wrapper for any of the cells
+        const wrapper =
+            notebook.cellCount > 0 ? this.getTextDocumentWrapper(notebook.getCells()[0].document) : undefined;
+        return wrapper?.handleRefresh(notebook) || [];
+    }
+
     public handleClose(cell: TextDocument) {
         const wrapper = this.getTextDocumentWrapper(cell);
         return wrapper?.handleClose(cell) || [];

src/notebookMiddlewareAddon.ts

@@ -29,7 +29,6 @@ import {
     LinkedEditingRanges,
     Location,
     NotebookDocument,
-    notebooks,
     Position,
     Position as VPosition,
     ProviderResult,
@@ -46,8 +45,7 @@ import {
     TextEdit,
     Uri,
     WorkspaceEdit,
-    workspace,
-    NotebookCellsChangeEvent
+    workspace
 } from 'vscode';
 import {
     ConfigurationParams,
@@ -55,6 +53,7 @@ import {
     DidChangeTextDocumentNotification,
     DidCloseTextDocumentNotification,
     DidOpenTextDocumentNotification,
+    ExecuteCommandSignature,
     HandleDiagnosticsSignature,
     LanguageClient,
     Middleware,
@@ -126,15 +125,13 @@ export class NotebookMiddlewareAddon implements Middleware, Disposable {
 
         // Make sure a bunch of functions are bound to this. VS code can call them without a this context
         this.handleDiagnostics = this.handleDiagnostics.bind(this);
+        this.executeCommand = this.executeCommand.bind(this);
         this.didOpen = this.didOpen.bind(this);
         this.didSave = this.didSave.bind(this);
         this.didChange = this.didChange.bind(this);
         this.didClose = this.didClose.bind(this);
         this.willSave = this.willSave.bind(this);
         this.willSaveWaitUntil = this.willSaveWaitUntil.bind(this);
-
-        // Since there are no LSP requests for moving cells around, we need to handle this ourselves
-        notebooks.onDidChangeNotebookCells(this.onDidChangeNotebookCells, this, this.disposables);
     }
 
     public workspace = {
@@ -168,6 +165,19 @@ export class NotebookMiddlewareAddon implements Middleware, Disposable {
         this.converter.dispose();
     }
 
+    public executeCommand(command: string, args: any[], next: ExecuteCommandSignature) {
+        const client = this.getClient();
+
+        // Pass onto the server unless this is our special command for handling cell movement (no API in LSP for this yet)
+        if (command === 'notebook.refresh' && client) {
+            // Send this to our converter and then the change notification to the server
+            const params = this.converter.handleRefresh(args[0]);
+            client.sendNotification(DidChangeTextDocumentNotification.type, params);
+        } else {
+            next(command, args);
+        }
+    }
+
     public stopWatching(notebook: NotebookDocument): void {
         // Just close the document. This should cause diags and other things to be cleared
         const client = this.getClient();
@@ -856,14 +866,6 @@ export class NotebookMiddlewareAddon implements Middleware, Disposable {
         }
     }
 
-    private onDidChangeNotebookCells(ev: NotebookCellsChangeEvent) {
-        // Each of these has to be turned into a change event so the underlying concat document
-        // is correct
-        if (this.shouldProvideIntellisense(ev.document.uri)) {
-            console.log(`Did change ${JSON.stringify(ev)}`);
-        }
-    }
-
     private shouldProvideIntellisense(uri: Uri): boolean {
         // Make sure document is allowed
         return this.isDocumentAllowed(uri);

src/notebookWrapper.ts

@@ -70,6 +70,26 @@ export class NotebookWrapper implements vscode.Disposable {
             return result;
         }
     }
+    public handleRefresh(notebook: vscode.NotebookDocument): protocol.DidChangeTextDocumentParams | undefined {
+        if (notebook == this.notebook) {
+            // Convert the notebook into something the concat document can understand (protocol types)
+            return this.concatDocument.handleRefresh({
+                cells: notebook
+                    .getCells()
+                    .filter((c) => score(c.document, this.selector))
+                    .map((c) => {
+                        return {
+                            textDocument: {
+                                uri: c.document.uri.toString(),
+                                version: c.document.version,
+                                languageId: c.document.languageId,
+                                text: c.document.getText()
+                            }
+                        };
+                    })
+            });
+        }
+    }
     public getText(range?: vscode.Range) {
         return this.concatDocument.getText(range);
     }

src/test/runTest.ts

@@ -66,7 +66,7 @@ async function main() {
             extensionTestsPath,
             launchArgs: [
                 workspacePath,
-                '--enable-proposed-api=vscode-jupyter-lsp-middleware',
+                '--enable-proposed-api=ms-vscode.vscode-jupyter-lsp-middleware',
                 '--skip-welcome',
                 '--skip-release-notes',
                 '--disable-workspace-trust'

src/test/suite/helper.ts

@@ -17,6 +17,7 @@ import {
     ExecuteCommandRegistrationOptions,
     ExecuteCommandRequest,
     LanguageClient,
+    Middleware,
     RegistrationData,
     RegistrationType,
     RevealOutputChannelOn,
@@ -873,6 +874,34 @@ function createMiddleware(
     }
 }
 
+function trackNotebookCellMovement(middleware: Middleware): vscode.Disposable {
+    return vscode.notebooks.onDidChangeNotebookCells((e) => {
+        // Translate notebook cell movement into change events
+        if (middleware.executeCommand) {
+            // If more than one item changed, then a move or a delete of multiple happened. Move
+            // always has a delete and a non delete
+            const deletions = e.changes
+                .filter((i) => i.deletedCount > 0)
+                .map((v) => {
+                    return { start: v.start, uris: v.deletedItems.map((d) => d.document.uri.toString()) };
+                });
+            const adds = e.changes
+                .filter((i) => i.deletedCount <= 0)
+                .map((v) => {
+                    return { start: v.start, uris: v.items.map((d) => d.document.uri.toString()) };
+                });
+            const isMove = adds.length > 0 && deletions.length == adds.length;
+
+            // If adds and deletes are same length, we should have done a move
+            if (isMove) {
+                middleware.executeCommand('notebook.refresh', [e.document], (_c, _args) => {
+                    // Do nothing as we don't need to send this anywhere
+                });
+            }
+        }
+    });
+}
+
 async function startLanguageServer(
     outputChannel: string,
     languageServerFolder: string,
@@ -914,6 +943,9 @@ async function startLanguageServer(
         shouldProvideIntellisense
     );
 
+    // Add custom message handling for notebook cell movement
+    const trackingDisposable = trackNotebookCellMovement(middleware);
+
     // Client options need to include our middleware piece
     const clientOptions: vslc.LanguageClientOptions = {
         documentSelector: selector,
@@ -948,7 +980,7 @@ async function startLanguageServer(
         await languageClient.onReady();
     }
 
-    return new LanguageServer(languageClient, [languageClientDisposable, cancellationStrategy]);
+    return new LanguageServer(languageClient, [languageClientDisposable, cancellationStrategy, trackingDisposable]);
 }
 
 export async function createLanguageServer(

src/test/suite/notebook.test.ts

@@ -209,6 +209,33 @@ suite('Notebook tests', function () {
             'System message not found'
         );
     });
+    test('Move cell with variable up and down (and make sure diags appear)', async () => {
+        await insertCodeCell('x = 4');
+        await insertMarkdownCell('# HEADER1');
+        const cell3 = await insertCodeCell('print(x)');
+
+        await focusCell(cell3);
+        let changePromise = waitForCellChange();
+        await commands.executeCommand('notebook.cell.moveUp');
+        await changePromise;
+
+        // Move again
+        const cell2 = window.activeNotebookEditor?.document.cellAt(1)!;
+
+        await focusCell(cell2);
+        changePromise = waitForCellChange();
+        await commands.executeCommand('notebook.cell.moveUp');
+        await changePromise;
+
+        // First cell should have diags now
+        const cell1 = window.activeNotebookEditor?.document.cellAt(0)!;
+        const diagnostics = await waitForDiagnostics(cell1.document.uri);
+        assert.ok(diagnostics, 'Moving variable down should cause diags to show up');
+        assert.ok(
+            diagnostics.find((item) => item.message.includes('x')),
+            'X message not found'
+        );
+    });
     test('Add some errors with markdown and delete some cells', async () => {
         await insertCodeCell('import sys\nprint(sys.executable)');
         await insertMarkdownCell('# HEADER1');