@marimo-team/frontend 0.23.7-dev5 → 0.23.7-dev51

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

@@ -101,6 +101,158 @@ print(x)`);
     `);
   });
 
+  test("selects the nearest in-scope local definition", async () => {
+    const code = `\
+a = 10
+
+def my_func():
+    a = 20
+    print(a)`;
+    view = createEditor(code);
+    const result = goToVariableDefinition(view, "a", code.lastIndexOf("a"));
+
+    expect(result).toBe(true);
+    await tick();
+    expect(renderEditorView(view)).toMatchInlineSnapshot(`
+      "
+      a = 10
+
+      def my_func():
+          a = 20
+          ^
+          print(a)
+      "
+    `);
+  });
+
+  test("selects the nearest in-scope parameter definition", async () => {
+    const code = `\
+a = 10
+
+def my_func(a):
+    print(a)`;
+    view = createEditor(code);
+    const result = goToVariableDefinition(view, "a", code.lastIndexOf("a"));
+
+    expect(result).toBe(true);
+    await tick();
+    expect(renderEditorView(view)).toMatchInlineSnapshot(`
+      "
+      a = 10
+
+      def my_func(a):
+                  ^
+          print(a)
+      "
+    `);
+  });
+
+  test("selects the comprehension target inside a set comprehension", async () => {
+    const code = `\
+x = 100
+s = {x for x in range(10)}`;
+    view = createEditor(code);
+    // Go-to-definition on the `x` before `for` (the expression part of the
+    // comprehension).
+    const usagePosition = code.indexOf("{x") + 1;
+    const result = goToVariableDefinition(view, "x", usagePosition);
+
+    expect(result).toBe(true);
+    await tick();
+    // Should jump to the comprehension target `x` (after `for`), not the
+    // outer `x = 100`. The Lezer Python grammar emits
+    // `SetComprehensionExpression`, and we now correctly match it in
+    // SCOPE_CREATING_NODES, so the comprehension creates a scope and the
+    // for-target is collected correctly.
+    expect(renderEditorView(view)).toMatchInlineSnapshot(`
+      "
+      x = 100
+      s = {x for x in range(10)}
+                 ^
+      "
+    `);
+  });
+
+  test("selects the comprehension target inside a dict comprehension", async () => {
+    const code = `\
+x = 100
+d = {x: x for x in range(10)}`;
+    view = createEditor(code);
+    const usagePosition = code.indexOf("{x") + 1;
+    const result = goToVariableDefinition(view, "x", usagePosition);
+
+    expect(result).toBe(true);
+    await tick();
+    // Positive control: `DictionaryComprehensionExpression` matches the grammar
+    // and is in SCOPE_CREATING_NODES, so this should jump to the comprehension
+    // target `x` (after `for`).
+    expect(renderEditorView(view)).toMatchInlineSnapshot(`
+      "
+      x = 100
+      d = {x: x for x in range(10)}
+                    ^
+      "
+    `);
+  });
+
+  test("skips enclosing class scope when resolving from inside a method", async () => {
+    const code = `\
+x = 100
+class Foo:
+    x = 10
+    def method(self):
+        return x`;
+    view = createEditor(code);
+    // Go-to-definition on the `x` inside `return x`.
+    const usagePosition = code.lastIndexOf("x");
+    const result = goToVariableDefinition(view, "x", usagePosition);
+
+    expect(result).toBe(true);
+    await tick();
+    // Should jump to `x = 100` at module scope. In Python, methods do NOT see
+    // their enclosing class body's names — class scopes are skipped in LEGB
+    // lookup once a function boundary has been crossed. We now correctly skip
+    // ClassDefinition in getScopeChain once a function boundary is crossed.
+    expect(renderEditorView(view)).toMatchInlineSnapshot(`
+      "
+      x = 100
+      ^
+      class Foo:
+          x = 10
+          def method(self):
+              return x
+      "
+    `);
+  });
+
+  test("resolves a global forward-reference from inside a function", async () => {
+    const code = `\
+def foo():
+    return a
+
+a = 10`;
+    view = createEditor(code);
+    // Go-to-definition on the `a` inside `return a`.
+    const usagePosition = code.indexOf("return a") + "return ".length;
+    const result = goToVariableDefinition(view, "a", usagePosition);
+
+    expect(result).toBe(true);
+    await tick();
+    // Should jump to `a = 10` at the bottom. Python allows forward references
+    // from within nested functions to module-level names. We now correctly omit
+    // "global" from POSITION_SENSITIVE_SCOPES, allowing forward references to
+    // global-level definitions declared after the usage position.
+    expect(renderEditorView(view)).toMatchInlineSnapshot(`
+      "
+      def foo():
+          return a
+
+      a = 10
+      ^
+      "
+    `);
+  });
+
   test("selects outer-scope function declaration", async () => {
     view = createEditor(`\
 def x():

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

@@ -0,0 +1,99 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { python } from "@codemirror/lang-python";
+import { EditorState } from "@codemirror/state";
+import { EditorView } from "@codemirror/view";
+import { afterEach, describe, expect, test } from "vitest";
+import { cellId, variableName } from "@/__tests__/branded";
+import { initialNotebookState, notebookAtom } from "@/core/cells/cells";
+import { store } from "@/core/state/jotai";
+import { variablesAtom } from "@/core/variables/state";
+import { goToDefinitionAtCursorPosition } from "../utils";
+
+async function tick(): Promise<void> {
+  await new Promise((resolve) => requestAnimationFrame(resolve));
+}
+
+function createEditor(content: string, selection: number) {
+  const state = EditorState.create({
+    doc: content,
+    selection: { anchor: selection },
+    extensions: [python()],
+  });
+
+  return new EditorView({
+    state,
+    parent: document.body,
+  });
+}
+
+const views: EditorView[] = [];
+
+afterEach(() => {
+  for (const view of views.splice(0)) {
+    view.destroy();
+  }
+
+  store.set(notebookAtom, initialNotebookState());
+  store.set(variablesAtom, {});
+});
+
+describe("goToDefinitionAtCursorPosition", () => {
+  test("prefers the current-cell local definition over a reactive global", async () => {
+    const globalCell = cellId("global-cell");
+    const localCell = cellId("local-cell");
+    const globalCode = `\
+a = 10
+print(a)`;
+    const localCode = `\
+def test():
+    a = 20
+    print(a)`;
+
+    const globalView = createEditor(globalCode, globalCode.length);
+    const localView = createEditor(localCode, localCode.lastIndexOf("a"));
+    views.push(globalView, localView);
+
+    const notebook = initialNotebookState();
+    notebook.cellHandles[globalCell] = {
+      current: { editorView: globalView, editorViewOrNull: globalView },
+    };
+    notebook.cellHandles[localCell] = {
+      current: { editorView: localView, editorViewOrNull: localView },
+    };
+
+    store.set(notebookAtom, notebook);
+    store.set(variablesAtom, {
+      [variableName("a")]: {
+        dataType: "int",
+        declaredBy: [globalCell],
+        name: variableName("a"),
+        usedBy: [localCell],
+        value: "10",
+      },
+    });
+
+    const result = goToDefinitionAtCursorPosition(localView);
+
+    expect(result).toBe(true);
+    await tick();
+    expect(localView.state.selection.main.head).toBe(
+      localCode.indexOf("a = 20"),
+    );
+    expect(globalView.state.selection.main.head).toBe(globalCode.length);
+  });
+
+  test("keeps private variables within the current cell", async () => {
+    const code = `\
+_x = 10
+output = _x + 10`;
+    const view = createEditor(code, code.lastIndexOf("_x"));
+    views.push(view);
+
+    const result = goToDefinitionAtCursorPosition(view);
+
+    expect(result).toBe(true);
+    await tick();
+    expect(view.state.selection.main.head).toBe(code.indexOf("_x = 10"));
+  });
+});

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

@@ -1,7 +1,31 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { syntaxTree } from "@codemirror/language";
+import type { EditorState } from "@codemirror/state";
 import { EditorView } from "@codemirror/view";
+import type { SyntaxNode, Tree, TreeCursor } from "@lezer/common";
+
+const SCOPE_CREATING_NODES = new Set([
+  "FunctionDefinition",
+  "LambdaExpression",
+  "ArrayComprehensionExpression",
+  "SetComprehensionExpression",
+  "DictionaryComprehensionExpression",
+  "ComprehensionExpression",
+  "ClassDefinition",
+]);
+
+const POSITION_SENSITIVE_SCOPES = new Set(["ClassDefinition"]);
+
+interface ScopeContext {
+  id: number;
+  type: string;
+}
+
+interface VariableDeclaration {
+  from: number;
+  scopeId: number;
+}
 
 function goToPosition(view: EditorView, from: number): void {
   view.focus();
@@ -22,50 +46,386 @@ function goToPosition(view: EditorView, from: number): void {
   });
 }
 
-/**
- * This function will select the first occurrence of the given variable name,
- * for a given editor view.
- * @param view The editor view which contains the variable name.
- * @param variableName The name of the variable to select, if found in the editor.
- */
-export function goToVariableDefinition(
-  view: EditorView,
+function findFirstMatchingVariable(
+  state: EditorState,
   variableName: string,
-): boolean {
-  const { state } = view;
+): number | null {
   const tree = syntaxTree(state);
 
-  let found = false;
-  let from = 0;
+  let from: number | null = null;
 
   tree.iterate({
     enter: (node) => {
-      if (found) {
+      if (from !== null) {
         return false;
-      } // Stop traversal if found
+      }
 
-      // Check if the node is an identifier and matches the variable name
       if (
         node.name === "VariableName" &&
         state.doc.sliceString(node.from, node.to) === variableName
       ) {
         from = node.from;
-        found = true;
-        return false; // Stop traversal
+        return false;
       }
 
-      // Skip comments and strings
       if (node.name === "Comment" || node.name === "String") {
         return false;
       }
+
+      return undefined;
     },
   });
 
-  if (found) {
-    goToPosition(view, from);
-    return true;
+  return from;
+}
+
+function getScopeChain(tree: Tree, usagePosition: number): ScopeContext[] {
+  const scopeChain: ScopeContext[] = [];
+  let currentNode: SyntaxNode | null = tree.resolveInner(usagePosition, 0);
+
+  while (currentNode) {
+    if (SCOPE_CREATING_NODES.has(currentNode.name)) {
+      // Skip ClassDefinition if we've already seen a function/lambda.
+      const inFunctionLikeScope = scopeChain.some(
+        (scope) =>
+          scope.type === "FunctionDefinition" ||
+          scope.type === "LambdaExpression",
+      );
+      if (!(inFunctionLikeScope && currentNode.name === "ClassDefinition")) {
+        scopeChain.push({
+          id: currentNode.from,
+          type: currentNode.name,
+        });
+      }
+    }
+    currentNode = currentNode.parent;
+  }
+
+  scopeChain.push({ id: -1, type: "global" });
+  return scopeChain;
+}
+
+function addDeclaration(
+  declarations: VariableDeclaration[],
+  scopeId: number,
+  from: number,
+) {
+  declarations.push({ scopeId, from });
+}
+
+function traverseChildren(
+  cursor: TreeCursor,
+  callback: (node: SyntaxNode) => void,
+) {
+  if (cursor.firstChild()) {
+    do {
+      callback(cursor.node);
+    } while (cursor.nextSibling());
+  }
+}
+
+function collectMatchingTargets(
+  cursor: TreeCursor,
+  state: EditorState,
+  variableName: string,
+  scopeId: number,
+  declarations: VariableDeclaration[],
+) {
+  switch (cursor.name) {
+    case "VariableName":
+      if (state.doc.sliceString(cursor.from, cursor.to) === variableName) {
+        addDeclaration(declarations, scopeId, cursor.from);
+      }
+      break;
+
+    case "TupleExpression":
+    case "ArrayExpression": {
+      const childCursor = cursor.node.cursor();
+      childCursor.firstChild();
+      do {
+        collectMatchingTargets(
+          childCursor,
+          state,
+          variableName,
+          scopeId,
+          declarations,
+        );
+      } while (childCursor.nextSibling());
+      break;
+    }
+    default:
+      break;
   }
-  return false;
+}
+
+function collectFunctionParameters(
+  node: SyntaxNode | Tree,
+  state: EditorState,
+  variableName: string,
+  scopeId: number,
+  declarations: VariableDeclaration[],
+) {
+  const cursor = node.cursor();
+  cursor.firstChild();
+  do {
+    if (cursor.name !== "ParamList") {
+      continue;
+    }
+
+    const paramCursor = cursor.node.cursor();
+    paramCursor.firstChild();
+    do {
+      if (
+        paramCursor.name === "VariableName" &&
+        state.doc.sliceString(paramCursor.from, paramCursor.to) === variableName
+      ) {
+        addDeclaration(declarations, scopeId, paramCursor.from);
+      }
+    } while (paramCursor.nextSibling());
+  } while (cursor.nextSibling());
+}
+
+function collectForTargets(
+  node: SyntaxNode | Tree,
+  state: EditorState,
+  variableName: string,
+  scopeId: number,
+  declarations: VariableDeclaration[],
+) {
+  const cursor = node.cursor();
+  cursor.firstChild();
+  let foundFor = false;
+  do {
+    if (cursor.name === "for") {
+      foundFor = true;
+    } else if (foundFor && cursor.name === "in") {
+      break;
+    } else if (foundFor) {
+      collectMatchingTargets(
+        cursor,
+        state,
+        variableName,
+        scopeId,
+        declarations,
+      );
+    }
+  } while (cursor.nextSibling());
+}
+
+function collectMatchingDeclarations(
+  node: SyntaxNode | Tree,
+  state: EditorState,
+  variableName: string,
+  scopeStack: number[],
+  declarations: VariableDeclaration[],
+) {
+  const cursor = node.cursor();
+  const nodeName = cursor.name;
+  const nodeStart = cursor.from;
+
+  const isNewScope = SCOPE_CREATING_NODES.has(nodeName);
+  const currentScopeStack = isNewScope
+    ? [...scopeStack, nodeStart]
+    : scopeStack;
+  const currentScope = currentScopeStack[currentScopeStack.length - 1] ?? -1;
+
+  switch (nodeName) {
+    case "FunctionDefinition":
+    case "ClassDefinition": {
+      const subCursor = node.cursor();
+      subCursor.firstChild();
+      do {
+        if (
+          subCursor.name === "VariableName" &&
+          state.doc.sliceString(subCursor.from, subCursor.to) === variableName
+        ) {
+          const parentScope = scopeStack[scopeStack.length - 1] ?? -1;
+          addDeclaration(declarations, parentScope, subCursor.from);
+          break;
+        }
+      } while (subCursor.nextSibling());
+
+      if (nodeName === "FunctionDefinition") {
+        collectFunctionParameters(
+          node,
+          state,
+          variableName,
+          nodeStart,
+          declarations,
+        );
+      }
+      break;
+    }
+    case "LambdaExpression":
+      collectFunctionParameters(
+        node,
+        state,
+        variableName,
+        nodeStart,
+        declarations,
+      );
+      break;
+
+    case "ArrayComprehensionExpression":
+    case "DictionaryComprehensionExpression":
+    case "SetComprehensionExpression":
+    case "ComprehensionExpression":
+    case "ForStatement":
+      collectForTargets(node, state, variableName, currentScope, declarations);
+      break;
+
+    case "AssignStatement": {
+      const assignOpPositions: number[] = [];
+      const subCursor = node.cursor();
+      subCursor.firstChild();
+      do {
+        if (subCursor.name === "AssignOp") {
+          assignOpPositions.push(subCursor.from);
+        }
+      } while (subCursor.nextSibling());
+
+      const lastAssignOpPosition =
+        assignOpPositions[assignOpPositions.length - 1];
+      if (lastAssignOpPosition === undefined) {
+        break;
+      }
+
+      const targetCursor = node.cursor();
+      targetCursor.firstChild();
+      do {
+        if (targetCursor.from < lastAssignOpPosition) {
+          collectMatchingTargets(
+            targetCursor,
+            state,
+            variableName,
+            currentScope,
+            declarations,
+          );
+        }
+      } while (targetCursor.nextSibling());
+      break;
+    }
+    case "ImportStatement": {
+      const subCursor = node.cursor();
+      subCursor.firstChild();
+      do {
+        if (
+          subCursor.name === "VariableName" &&
+          state.doc.sliceString(subCursor.from, subCursor.to) === variableName
+        ) {
+          addDeclaration(declarations, currentScope, subCursor.from);
+        }
+      } while (subCursor.nextSibling());
+      break;
+    }
+    case "ImportFromStatement": {
+      const subCursor = node.cursor();
+      subCursor.firstChild();
+      let foundImport = false;
+      do {
+        if (subCursor.name === "import") {
+          foundImport = true;
+        } else if (
+          foundImport &&
+          subCursor.name === "VariableName" &&
+          state.doc.sliceString(subCursor.from, subCursor.to) === variableName
+        ) {
+          addDeclaration(declarations, currentScope, subCursor.from);
+        }
+      } while (subCursor.nextSibling());
+      break;
+    }
+    case "TryStatement":
+    case "WithStatement": {
+      const subCursor = node.cursor();
+      subCursor.firstChild();
+      let foundAs = false;
+      do {
+        if (subCursor.name === "as") {
+          foundAs = true;
+        } else if (
+          foundAs &&
+          subCursor.name === "VariableName" &&
+          state.doc.sliceString(subCursor.from, subCursor.to) === variableName
+        ) {
+          addDeclaration(declarations, currentScope, subCursor.from);
+          foundAs = false;
+        }
+      } while (subCursor.nextSibling());
+      break;
+    }
+    default:
+      break;
+  }
+
+  traverseChildren(cursor, (childNode) => {
+    collectMatchingDeclarations(
+      childNode,
+      state,
+      variableName,
+      currentScopeStack,
+      declarations,
+    );
+  });
+}
+
+function findScopedDefinitionPosition(
+  state: EditorState,
+  variableName: string,
+  usagePosition: number,
+): number | null {
+  const tree = syntaxTree(state);
+  const declarations: VariableDeclaration[] = [];
+
+  collectMatchingDeclarations(tree, state, variableName, [], declarations);
+
+  const clampedUsagePosition = Math.max(
+    0,
+    Math.min(usagePosition, state.doc.length),
+  );
+
+  for (const scope of getScopeChain(tree, clampedUsagePosition)) {
+    const match = declarations
+      .filter((declaration) => declaration.scopeId === scope.id)
+      .filter((declaration) => {
+        return POSITION_SENSITIVE_SCOPES.has(scope.type)
+          ? declaration.from <= clampedUsagePosition
+          : true;
+      })
+      .toSorted((left, right) => left.from - right.from)[0];
+
+    if (match) {
+      return match.from;
+    }
+  }
+
+  return null;
+}
+
+/**
+ * This function will select the first occurrence of the given variable name,
+ * for a given editor view.
+ * @param view The editor view which contains the variable name.
+ * @param variableName The name of the variable to select, if found in the editor.
+ * @param usagePosition The position of the variable usage, if available.
+ */
+export function goToVariableDefinition(
+  view: EditorView,
+  variableName: string,
+  usagePosition?: number,
+): boolean {
+  const { state } = view;
+  const from =
+    (usagePosition !== undefined
+      ? findScopedDefinitionPosition(state, variableName, usagePosition)
+      : null) ?? findFirstMatchingVariable(state, variableName);
+
+  if (from === null) {
+    return false;
+  }
+
+  goToPosition(view, from);
+  return true;
 }
 
 /**