@marimo-team/islands 0.23.16-dev9 → 0.23.17-dev0

package/src/core/codemirror/find-replace/__tests__/navigate.test.ts

@@ -20,6 +20,26 @@ vi.mock("@/core/cells/cells", () => ({
   getAllEditorViews: () => mockGetAllEditorViews(),
 }));
 
+/**
+ * A view stub that only carries what navigate needs
+ */
+function createMockView(state: EditorState): EditorView {
+  const cell = document.createElement("div");
+  cell.className = "marimo-cell";
+  // jsdom doesn't implement scrollIntoView.
+  cell.scrollIntoView = vi.fn();
+  const dom = document.createElement("div");
+  cell.append(dom);
+
+  return {
+    state,
+    dom,
+    // Don't actually apply changes; it's complex to mock properly and the
+    // assertions only care that dispatch was called.
+    dispatch: vi.fn(),
+  } as unknown as EditorView;
+}
+
 describe("navigate", () => {
   let view1: EditorView;
   let view2: EditorView;
@@ -36,20 +56,8 @@ describe("navigate", () => {
     });
 
     // Create mock views that track dispatch calls
-    view1 = {
-      state: state1,
-      dispatch: vi.fn(() => {
-        // Mock dispatch - in real tests we just need to verify it was called
-        // Don't actually apply changes as it's complex to mock properly
-      }),
-    } as unknown as EditorView;
-
-    view2 = {
-      state: state2,
-      dispatch: vi.fn(() => {
-        // Mock dispatch - in real tests we just need to verify it was called
-      }),
-    } as unknown as EditorView;
+    view1 = createMockView(state1);
+    view2 = createMockView(state2);
 
     mockViews = [view1, view2];
     mockGetAllEditorViews.mockReturnValue(mockViews);
@@ -234,6 +242,22 @@ describe("navigate", () => {
 
       expect(result).toBe(false);
     });
+
+    it("should scroll the cell owning the match into view", async () => {
+      const cell = view1.dom.closest(".marimo-cell");
+      invariant(cell, "mock view should be inside a cell");
+
+      expect(findNext()).toBeTruthy();
+
+      // The scroll is deferred a frame.
+      await new Promise((resolve) => requestAnimationFrame(resolve));
+
+      expect(cell.scrollIntoView).toHaveBeenCalledWith({
+        behavior: "instant",
+        block: "nearest",
+        inline: "nearest",
+      });
+    });
   });
 
   describe("findPrev", () => {
@@ -319,10 +343,7 @@ describe("navigate", () => {
         extensions: [EditorState.readOnly.of(true)],
       });
 
-      const readOnlyView = {
-        state: readOnlyState,
-        dispatch: vi.fn(),
-      } as unknown as EditorView;
+      const readOnlyView = createMockView(readOnlyState);
 
       // Replace view1 with read-only view temporarily
       mockGetAllEditorViews.mockReturnValueOnce([readOnlyView, view2]);
@@ -568,10 +589,7 @@ describe("navigate", () => {
         doc: Text.of(["hello world hello there hello"]),
       });
 
-      const testView = {
-        state: testDoc,
-        dispatch: vi.fn(),
-      } as unknown as EditorView;
+      const testView = createMockView(testDoc);
 
       mockGetAllEditorViews.mockReturnValue([testView]);
 
@@ -613,10 +631,7 @@ describe("navigate", () => {
         doc: Text.of([""]),
       });
 
-      const emptyView = {
-        state: emptyState,
-        dispatch: vi.fn(),
-      } as unknown as EditorView;
+      const emptyView = createMockView(emptyState);
 
       mockGetAllEditorViews.mockReturnValue([emptyView]);
 
@@ -629,10 +644,7 @@ describe("navigate", () => {
         doc: Text.of(["orphan content"]),
       });
 
-      const orphanView = {
-        state: orphanState,
-        dispatch: vi.fn(),
-      } as unknown as EditorView;
+      const orphanView = createMockView(orphanState);
 
       store.set(findReplaceAtom, {
         type: "setCurrentView",

package/src/core/codemirror/find-replace/navigate.ts

@@ -6,6 +6,7 @@ import { EditorView } from "@codemirror/view";
 import { getAllEditorViews } from "@/core/cells/cells";
 import { replaceEditorContent } from "@/core/codemirror/replace-editor-content";
 import { store } from "@/core/state/jotai";
+import { scrollOwnerCell } from "../utils";
 import { asQueryCreator, type QueryType } from "./query";
 import { findReplaceAtom } from "./state";
 
@@ -85,6 +86,7 @@ const findInDirection = (direction: "next" | "prev") =>
         effects: [EditorView.scrollIntoView(selection.main, { y: "center" })],
         userEvent: "select.search",
       });
+      scrollOwnerCell(view);
       store.set(findReplaceAtom, {
         type: "setCurrentView",
         view,

package/src/core/codemirror/go-to-definition/__tests__/commands.test.ts

@@ -3,14 +3,14 @@
 import { python } from "@codemirror/lang-python";
 import { EditorState } from "@codemirror/state";
 import { EditorView } from "@codemirror/view";
-import { afterEach, describe, expect, test } from "vitest";
+import { afterEach, describe, expect, test, vi } from "vitest";
 import { goToVariableDefinition } from "../commands";
 
 async function tick(): Promise<void> {
   await new Promise((resolve) => requestAnimationFrame(resolve));
 }
 
-function createEditor(content: string) {
+function createEditor(content: string, parent: HTMLElement = document.body) {
   const state = EditorState.create({
     doc: content,
     extensions: [python()],
@@ -18,7 +18,7 @@ function createEditor(content: string) {
 
   const view = new EditorView({
     state,
-    parent: document.body,
+    parent,
   });
 
   return view;
@@ -383,6 +383,27 @@ print('myVar')`);
       "
     `);
   });
+
+  test("scrolls the cell owning the definition into view", async () => {
+    const cell = document.createElement("div");
+    cell.className = "marimo-cell";
+    // jsdom doesn't implement scrollIntoView.
+    cell.scrollIntoView = vi.fn();
+    document.body.append(cell);
+
+    view = createEditor("myVar = 10\nprint(myVar)", cell);
+    expect(goToVariableDefinition(view, "myVar")).toBe(true);
+    // The jump and the scroll are both deferred to the next frame.
+    await tick();
+
+    expect(cell.scrollIntoView).toHaveBeenCalledWith({
+      behavior: "instant",
+      block: "nearest",
+      inline: "nearest",
+    });
+
+    cell.remove();
+  });
 });
 
 /**

package/src/core/codemirror/go-to-definition/commands.ts

@@ -4,19 +4,13 @@ import { syntaxTree } from "@codemirror/language";
 import type { EditorState } from "@codemirror/state";
 import { EditorView } from "@codemirror/view";
 import type { SyntaxNode, Tree, TreeCursor } from "@lezer/common";
+import { PyKeyword, PyNode, SCOPE_CREATING_NODES } from "../python-node-names";
+import { scrollOwnerCell } from "../utils";
 
-const SCOPE_CREATING_NODES = new Set([
-  "FunctionDefinition",
-  "LambdaExpression",
-  "ArrayComprehensionExpression",
-  "SetComprehensionExpression",
-  "DictionaryComprehensionExpression",
-  "ComprehensionExpression",
-  "ClassDefinition",
+const POSITION_SENSITIVE_SCOPES: ReadonlySet<string> = new Set([
+  PyNode.ClassDefinition,
 ]);
 
-const POSITION_SENSITIVE_SCOPES = new Set(["ClassDefinition"]);
-
 interface ScopeContext {
   id: number;
   type: string;
@@ -45,6 +39,8 @@ function goToPosition(view: EditorView, from: number): void {
       }),
     });
   });
+
+  scrollOwnerCell(view);
 }
 
 function findFirstMatchingVariable(
@@ -62,14 +58,14 @@ function findFirstMatchingVariable(
       }
 
       if (
-        node.name === "VariableName" &&
+        node.name === PyNode.VariableName &&
         state.doc.sliceString(node.from, node.to) === variableName
       ) {
         from = node.from;
         return false;
       }
 
-      if (node.name === "Comment" || node.name === "String") {
+      if (node.name === PyNode.Comment || node.name === PyNode.String) {
         return false;
       }
 
@@ -89,10 +85,12 @@ function getScopeChain(tree: Tree, usagePosition: number): ScopeContext[] {
       // Skip ClassDefinition if we've already seen a function/lambda.
       const inFunctionLikeScope = scopeChain.some(
         (scope) =>
-          scope.type === "FunctionDefinition" ||
-          scope.type === "LambdaExpression",
+          scope.type === PyNode.FunctionDefinition ||
+          scope.type === PyNode.LambdaExpression,
       );
-      if (!(inFunctionLikeScope && currentNode.name === "ClassDefinition")) {
+      if (
+        !(inFunctionLikeScope && currentNode.name === PyNode.ClassDefinition)
+      ) {
         scopeChain.push({
           id: currentNode.from,
           type: currentNode.name,
@@ -127,30 +125,27 @@ function traverseChildren(
 
 function collectMatchingTargets(
   cursor: TreeCursor,
-  state: EditorState,
-  variableName: string,
-  scopeId: number,
-  declarations: VariableDeclaration[],
+  options: {
+    state: EditorState;
+    variableName: string;
+    scopeId: number;
+    declarations: VariableDeclaration[];
+  },
 ) {
+  const { state, variableName, scopeId, declarations } = options;
   switch (cursor.name) {
-    case "VariableName":
+    case PyNode.VariableName:
       if (state.doc.sliceString(cursor.from, cursor.to) === variableName) {
         addDeclaration(declarations, scopeId, cursor.from);
       }
       break;
 
-    case "TupleExpression":
-    case "ArrayExpression": {
+    case PyNode.TupleExpression:
+    case PyNode.ArrayExpression: {
       const childCursor = cursor.node.cursor();
       childCursor.firstChild();
       do {
-        collectMatchingTargets(
-          childCursor,
-          state,
-          variableName,
-          scopeId,
-          declarations,
-        );
+        collectMatchingTargets(childCursor, options);
       } while (childCursor.nextSibling());
       break;
     }
@@ -161,15 +156,18 @@ function collectMatchingTargets(
 
 function collectFunctionParameters(
   node: SyntaxNode | Tree,
-  state: EditorState,
-  variableName: string,
-  scopeId: number,
-  declarations: VariableDeclaration[],
+  options: {
+    state: EditorState;
+    variableName: string;
+    scopeId: number;
+    declarations: VariableDeclaration[];
+  },
 ) {
+  const { state, variableName, scopeId, declarations } = options;
   const cursor = node.cursor();
   cursor.firstChild();
   do {
-    if (cursor.name !== "ParamList") {
+    if (cursor.name !== PyNode.ParamList) {
       continue;
     }
 
@@ -177,7 +175,7 @@ function collectFunctionParameters(
     paramCursor.firstChild();
     do {
       if (
-        paramCursor.name === "VariableName" &&
+        paramCursor.name === PyNode.VariableName &&
         state.doc.sliceString(paramCursor.from, paramCursor.to) === variableName
       ) {
         addDeclaration(declarations, scopeId, paramCursor.from);
@@ -188,38 +186,37 @@ function collectFunctionParameters(
 
 function collectForTargets(
   node: SyntaxNode | Tree,
-  state: EditorState,
-  variableName: string,
-  scopeId: number,
-  declarations: VariableDeclaration[],
+  options: {
+    state: EditorState;
+    variableName: string;
+    scopeId: number;
+    declarations: VariableDeclaration[];
+  },
 ) {
   const cursor = node.cursor();
   cursor.firstChild();
   let foundFor = false;
   do {
-    if (cursor.name === "for") {
+    if (cursor.name === PyKeyword.For) {
       foundFor = true;
-    } else if (foundFor && cursor.name === "in") {
+    } else if (foundFor && cursor.name === PyKeyword.In) {
       break;
     } else if (foundFor) {
-      collectMatchingTargets(
-        cursor,
-        state,
-        variableName,
-        scopeId,
-        declarations,
-      );
+      collectMatchingTargets(cursor, options);
     }
   } while (cursor.nextSibling());
 }
 
 function collectMatchingDeclarations(
   node: SyntaxNode | Tree,
-  state: EditorState,
-  variableName: string,
-  scopeStack: number[],
-  declarations: VariableDeclaration[],
+  options: {
+    state: EditorState;
+    variableName: string;
+    scopeStack: number[];
+    declarations: VariableDeclaration[];
+  },
 ) {
+  const { state, variableName, scopeStack, declarations } = options;
   const cursor = node.cursor();
   const nodeName = cursor.name;
   const nodeStart = cursor.from;
@@ -231,13 +228,13 @@ function collectMatchingDeclarations(
   const currentScope = currentScopeStack[currentScopeStack.length - 1] ?? -1;
 
   switch (nodeName) {
-    case "FunctionDefinition":
-    case "ClassDefinition": {
+    case PyNode.FunctionDefinition:
+    case PyNode.ClassDefinition: {
       const subCursor = node.cursor();
       subCursor.firstChild();
       do {
         if (
-          subCursor.name === "VariableName" &&
+          subCursor.name === PyNode.VariableName &&
           state.doc.sliceString(subCursor.from, subCursor.to) === variableName
         ) {
           const parentScope = scopeStack[scopeStack.length - 1] ?? -1;
@@ -246,41 +243,44 @@ function collectMatchingDeclarations(
         }
       } while (subCursor.nextSibling());
 
-      if (nodeName === "FunctionDefinition") {
-        collectFunctionParameters(
-          node,
+      if (nodeName === PyNode.FunctionDefinition) {
+        collectFunctionParameters(node, {
           state,
           variableName,
-          nodeStart,
+          scopeId: nodeStart,
           declarations,
-        );
+        });
       }
       break;
     }
-    case "LambdaExpression":
-      collectFunctionParameters(
-        node,
+    case PyNode.LambdaExpression:
+      collectFunctionParameters(node, {
         state,
         variableName,
-        nodeStart,
+        scopeId: nodeStart,
         declarations,
-      );
+      });
       break;
 
-    case "ArrayComprehensionExpression":
-    case "DictionaryComprehensionExpression":
-    case "SetComprehensionExpression":
-    case "ComprehensionExpression":
-    case "ForStatement":
-      collectForTargets(node, state, variableName, currentScope, declarations);
+    case PyNode.ArrayComprehensionExpression:
+    case PyNode.DictionaryComprehensionExpression:
+    case PyNode.SetComprehensionExpression:
+    case PyNode.ComprehensionExpression:
+    case PyNode.ForStatement:
+      collectForTargets(node, {
+        state,
+        variableName,
+        scopeId: currentScope,
+        declarations,
+      });
       break;
 
-    case "AssignStatement": {
+    case PyNode.AssignStatement: {
       const assignOpPositions: number[] = [];
       const subCursor = node.cursor();
       subCursor.firstChild();
       do {
-        if (subCursor.name === "AssignOp") {
+        if (subCursor.name === PyNode.AssignOp) {
           assignOpPositions.push(subCursor.from);
         }
       } while (subCursor.nextSibling());
@@ -295,18 +295,17 @@ function collectMatchingDeclarations(
       targetCursor.firstChild();
       do {
         if (targetCursor.from < lastAssignOpPosition) {
-          collectMatchingTargets(
-            targetCursor,
+          collectMatchingTargets(targetCursor, {
             state,
             variableName,
-            currentScope,
+            scopeId: currentScope,
             declarations,
-          );
+          });
         }
       } while (targetCursor.nextSibling());
       break;
     }
-    case "ImportStatement": {
+    case PyNode.ImportStatement: {
       // The grammar emits one ImportStatement for both `import x [as y]` and
       // `from m import x [as y], ...`. Direct children include the keywords
       // (`from`/`import`/`as`), commas, dots, and every VariableName from the
@@ -327,19 +326,19 @@ function collectMatchingDeclarations(
         pending = null;
       };
       do {
-        if (subCursor.name === "import") {
+        if (subCursor.name === PyKeyword.Import) {
           pastImport = true;
           continue;
         }
         if (!pastImport) {
           continue;
         }
-        if (subCursor.name === "as") {
+        if (subCursor.name === PyKeyword.As) {
           // Next VariableName is the alias and replaces `pending`.
           pending = null;
           continue;
         }
-        if (subCursor.name === "VariableName") {
+        if (subCursor.name === PyNode.VariableName) {
           // Flush any previous pending name (no `as` followed it).
           commit();
           pending = {
@@ -355,17 +354,17 @@ function collectMatchingDeclarations(
       commit();
       break;
     }
-    case "TryStatement":
-    case "WithStatement": {
+    case PyNode.TryStatement:
+    case PyNode.WithStatement: {
       const subCursor = node.cursor();
       subCursor.firstChild();
       let foundAs = false;
       do {
-        if (subCursor.name === "as") {
+        if (subCursor.name === PyKeyword.As) {
           foundAs = true;
         } else if (
           foundAs &&
-          subCursor.name === "VariableName" &&
+          subCursor.name === PyNode.VariableName &&
           state.doc.sliceString(subCursor.from, subCursor.to) === variableName
         ) {
           addDeclaration(declarations, currentScope, subCursor.from);
@@ -379,13 +378,12 @@ function collectMatchingDeclarations(
   }
 
   traverseChildren(cursor, (childNode) => {
-    collectMatchingDeclarations(
-      childNode,
+    collectMatchingDeclarations(childNode, {
       state,
       variableName,
-      currentScopeStack,
+      scopeStack: currentScopeStack,
       declarations,
-    );
+    });
   });
 }
 
@@ -397,7 +395,12 @@ function findScopedDefinitionPosition(
   const tree = syntaxTree(state);
   const declarations: VariableDeclaration[] = [];
 
-  collectMatchingDeclarations(tree, state, variableName, [], declarations);
+  collectMatchingDeclarations(tree, {
+    state,
+    variableName,
+    scopeStack: [],
+    declarations,
+  });
 
   const clampedUsagePosition = Math.max(
     0,