diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 5436479342f..435ecfc5559 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -47,7 +47,7 @@ If an issue does not have one of these two labels, assume it is **not ready** fo ### Why is maintainer approval required? **Deliberate design.** marimo is an intentionally designed project. We -put just as much thought into the features we exclude as the ones we include, +put just as much thought into the features we exclude as the ones we include, in order to provide our users with a simple, consistent, delightful, and powerful experience. @@ -294,7 +294,7 @@ For the frontend, you can choose to run slower hot reloading for an environment - + diff --git a/frontend/src/core/codemirror/go-to-definition/__tests__/commands.test.ts b/frontend/src/core/codemirror/go-to-definition/__tests__/commands.test.ts index 2ca5447667b..85e7ef0e5f9 100644 --- a/frontend/src/core/codemirror/go-to-definition/__tests__/commands.test.ts +++ b/frontend/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(): diff --git a/frontend/src/core/codemirror/go-to-definition/__tests__/utils.test.ts b/frontend/src/core/codemirror/go-to-definition/__tests__/utils.test.ts new file mode 100644 index 00000000000..63e11cc2e26 --- /dev/null +++ b/frontend/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 { + 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")); + }); +}); diff --git a/frontend/src/core/codemirror/go-to-definition/commands.ts b/frontend/src/core/codemirror/go-to-definition/commands.ts index b67a193fc44..5ba4f47bd33 100644 --- a/frontend/src/core/codemirror/go-to-definition/commands.ts +++ b/frontend/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; } /** diff --git a/frontend/src/core/codemirror/go-to-definition/utils.ts b/frontend/src/core/codemirror/go-to-definition/utils.ts index 65db962881a..8eb09b7a413 100644 --- a/frontend/src/core/codemirror/go-to-definition/utils.ts +++ b/frontend/src/core/codemirror/go-to-definition/utils.ts @@ -18,10 +18,16 @@ function getWordUnderCursor(state: EditorState) { const { from, to } = state.selection.main; if (from === to) { const { startToken, endToken } = getPositionAtWordBounds(state.doc, from); - return state.doc.sliceString(startToken, endToken); + return { + position: startToken, + word: state.doc.sliceString(startToken, endToken), + }; } - return state.doc.sliceString(from, to); + return { + position: from, + word: state.doc.sliceString(from, to), + }; } /** @@ -51,15 +57,15 @@ function isPrivateVariable(variableName: string) { */ export function goToDefinitionAtCursorPosition(view: EditorView): boolean { const { state } = view; - const variableName = getWordUnderCursor(state); - if (!variableName) { + const { position, word } = getWordUnderCursor(state); + if (!word) { return false; } // Close popovers/tooltips closeCompletion(view); view.dispatch({ effects: closeHoverTooltips }); - return goToDefinition(view, variableName); + return goToDefinition(view, word, position); } /** @@ -69,7 +75,19 @@ export function goToDefinitionAtCursorPosition(view: EditorView): boolean { export function goToDefinition( view: EditorView, variableName: string, + usagePosition?: number, ): boolean { + if (usagePosition !== undefined) { + const foundLocally = goToVariableDefinition( + view, + variableName, + usagePosition, + ); + if (foundLocally) { + return true; + } + } + // The variable may exist in another cell const editorWithVariable = getEditorForVariable(view, variableName); if (!editorWithVariable) {
ProductionProduction-esque Development