Support scrolling in vscode notebooks

packages/app-vscode/src/ide/vscode/VscodeFocusEditor.ts

@@ -1,7 +1,10 @@
-import type { NotebookDocument, TextEditor } from "vscode";
+import type { TextEditor } from "vscode";
 import { commands, TabInputTextDiff, ViewColumn, window } from "vscode";
-import { getCellIndex } from "@cursorless/lib-vscode-common";
-import { getNotebookFromCellDocument } from "./notebook/notebook";
+import {
+  getCellIndex,
+  getNotebookFromCellDocument,
+  isDocumentNotebookCell,
+} from "@cursorless/lib-vscode-common";
 import type { VscodeTextEditor } from "./VscodeTextEditor";
 
 const columnFocusCommands = {
@@ -24,36 +27,36 @@ const columnFocusCommands = {
 export async function vscodeFocusEditor(
   editor: VscodeTextEditor,
 ): Promise<void> {
+  if (isDocumentNotebookCell(editor.vscodeEditor.document)) {
+    await focusNotebookCell(editor);
+    return;
+  }
+
   const viewColumn = getViewColumn(editor.vscodeEditor);
-  if (viewColumn != null) {
-    await commands.executeCommand(columnFocusCommands[viewColumn]);
 
-    if (editor.isDiffEditorOriginal && !editor.isActive) {
-      // There is no way of directly focusing the left hand side of a diff
-      // editor. Switch side if needed.
+  await commands.executeCommand(columnFocusCommands[viewColumn]);
 
-      await commands.executeCommand("diffEditor.switchSide");
-    } else if (editor.isSearchEditor) {
-      // Focusing the search editor brings focus back to the input field. This
-      // command moves selection into the actual editor text.
+  if (editor.isDiffEditorOriginal && !editor.isActive) {
+    // There is no way of directly focusing the left hand side of a diff
+    // editor. Switch side if needed.
 
-      await commands.executeCommand("search.action.focusNextSearchResult");
-    }
-  } else {
-    // If the view column is null we see if it's a notebook and try to see if we
-    // can just move around in the notebook to focus the correct editor
+    await commands.executeCommand("diffEditor.switchSide");
+  } else if (editor.isSearchEditor) {
+    // Focusing the search editor brings focus back to the input field. This
+    // command moves selection into the actual editor text.
 
-    await focusNotebookCell(editor);
+    await commands.executeCommand("search.action.focusNextSearchResult");
   }
 }
 
-function getViewColumn(editor: TextEditor): ViewColumn | undefined {
+function getViewColumn(editor: TextEditor): ViewColumn {
   if (editor.viewColumn != null) {
     return editor.viewColumn;
   }
+
   const uri = editor.document.uri.toString();
   const tabGroup = window.tabGroups.all.find((tabGroup) =>
-    tabGroup.tabs.find((tab) => {
+    tabGroup.tabs.some((tab) => {
       if (tab.input instanceof TabInputTextDiff) {
         return (
           tab.input.original.toString() === uri ||
@@ -63,41 +66,36 @@ function getViewColumn(editor: TextEditor): ViewColumn | undefined {
       return (tab as any).input?.uri?.toString() === uri;
     }),
   );
-  return tabGroup?.viewColumn;
+
+  const column = tabGroup?.viewColumn;
+
+  if (column == null) {
+    throw new Error("Couldn't determine view column for editor");
+  }
+
+  return column;
 }
 
 async function focusNotebookCell(editor: VscodeTextEditor) {
-  const desiredNotebookEditor = getNotebookFromCellDocument(
+  const notebookEditor = getNotebookFromCellDocument(
     editor.vscodeEditor.document,
   );
-  if (desiredNotebookEditor == null) {
-    throw new Error("Couldn't find notebook editor for given document");
-  }
-
-  const desiredNotebookDocument: NotebookDocument =
-    desiredNotebookEditor.notebook;
 
   await commands.executeCommand(
     columnFocusCommands[
-      desiredNotebookEditor.viewColumn as keyof typeof columnFocusCommands
+      notebookEditor.viewColumn as keyof typeof columnFocusCommands
     ],
   );
 
-  const desiredEditorIndex = getCellIndex(
-    desiredNotebookDocument,
+  const cellIndex = getCellIndex(
+    notebookEditor.notebook,
     editor.vscodeEditor.document,
   );
 
-  if (desiredEditorIndex < 0) {
-    throw new Error(
-      "Couldn't find cell corresponding to given document in notebook editor",
-    );
-  }
-
-  desiredNotebookEditor.selections = [
-    desiredNotebookEditor.selection.with({
-      start: desiredEditorIndex,
-      end: desiredEditorIndex + 1,
+  notebookEditor.selections = [
+    notebookEditor.selection.with({
+      start: cellIndex,
+      end: cellIndex + 1,
     }),
   ];
 

packages/app-vscode/src/ide/vscode/VscodeRevealLine.ts

@@ -1,12 +1,22 @@
 import vscode from "vscode";
 import { RevealLineAt } from "@cursorless/lib-common";
+import {
+  getCellIndex,
+  getNotebookFromCellDocument,
+  isDocumentNotebookCell,
+} from "@cursorless/lib-vscode-common";
 import type { VscodeTextEditor } from "./VscodeTextEditor";
 
 export async function vscodeRevealLine(
   editor: VscodeTextEditor,
   lineNumber: number,
   at: RevealLineAt,
 ): Promise<void> {
+  if (isDocumentNotebookCell(editor.vscodeEditor.document)) {
+    revealNotebookCell(editor, at);
+    return;
+  }
+
   // For reveal line to the work we have to have the correct editor focused
   if (!editor.isActive) {
     await editor.focus();
@@ -17,3 +27,38 @@ export async function vscodeRevealLine(
     at,
   });
 }
+
+function revealNotebookCell(editor: VscodeTextEditor, at: RevealLineAt) {
+  const notebookEditor = getNotebookFromCellDocument(
+    editor.vscodeEditor.document,
+  );
+
+  const cellIndex = getCellIndex(
+    notebookEditor.notebook,
+    editor.vscodeEditor.document,
+  );
+
+  notebookEditor.revealRange(
+    new vscode.NotebookRange(cellIndex, cellIndex),
+    getNotebookRevealType(at),
+  );
+}
+
+function getNotebookRevealType(
+  at: RevealLineAt,
+): vscode.NotebookEditorRevealType {
+  switch (at) {
+    case RevealLineAt.top:
+      return vscode.NotebookEditorRevealType.AtTop;
+    case RevealLineAt.center:
+      return vscode.NotebookEditorRevealType.InCenter;
+    case RevealLineAt.bottom:
+      throw new Error(
+        "VS Code's notebook API doesn't support bottom alignment.",
+      );
+    default: {
+      const _exhaustiveCheck: never = at;
+      throw new Error(`Unhandled case: ${_exhaustiveCheck}`);
+    }
+  }
+}

packages/lib-vscode-common/src/notebook.ts

@@ -1,4 +1,14 @@
-import type { NotebookDocument, TextDocument } from "vscode";
+import type { NotebookDocument, NotebookEditor, TextDocument } from "vscode";
+import { window } from "vscode";
+
+/**
+ * Checks if the given document corresponds to a notebook cell
+ * @param document The document to check
+ * @returns True if the document corresponds to a notebook cell, false otherwise
+ */
+export function isDocumentNotebookCell(document: TextDocument) {
+  return document.uri.scheme === "vscode-notebook-cell";
+}
 
 /**
  * Returns the index of the cell corresponding to the given document in the
@@ -11,10 +21,41 @@ import type { NotebookDocument, TextDocument } from "vscode";
 export function getCellIndex(
   notebookDocument: NotebookDocument,
   document: TextDocument,
-) {
-  return notebookDocument
+): number {
+  const index = notebookDocument
     .getCells()
     .findIndex(
       (cell) => cell.document.uri.toString() === document.uri.toString(),
     );
+
+  if (index === -1) {
+    throw new Error(
+      "Couldn't find cell corresponding to given document in notebook editor",
+    );
+  }
+
+  return index;
+}
+
+/**
+ * Given a document corresponding to a single cell, retrieve the notebook
+ * editor for the entire notebook
+ * @param document The document corresponding to the given cell
+ * @returns The notebook editor corresponding to the notebook containing the
+ * given cell
+ */
+export function getNotebookFromCellDocument(
+  document: TextDocument,
+): NotebookEditor {
+  const editor = window.visibleNotebookEditors.find((editor) =>
+    editor.notebook
+      .getCells()
+      .some((cell) => cell.document.uri.toString() === document.uri.toString()),
+  );
+
+  if (editor == null) {
+    throw new Error("Couldn't find notebook editor for given document");
+  }
+
+  return editor;
 }