Fix cross-file go-to-definition navigation in Pyrefly Sandbox - Issue #1061

pyrefly/lib/playground.rs

@@ -162,6 +162,12 @@ pub struct InlayHint {
     position: Position,
 }
 
+#[derive(Serialize)]
+pub struct DefinitionResult {
+    pub range: Range,
+    pub filename: String,
+}
+
 pub struct Playground {
     state: State,
     handles: HashMap<String, Handle>,
@@ -337,7 +343,7 @@ impl Playground {
         })
     }
 
-    pub fn goto_definition(&mut self, pos: Position) -> Option<Range> {
+    pub fn goto_definition(&mut self, pos: Position) -> Option<DefinitionResult> {
         let handle = self.handles.get(&self.active_filename)?;
         let transaction = self.state.transaction();
         let position = self.to_text_size(&transaction, pos)?;
@@ -346,7 +352,20 @@ impl Playground {
             .goto_definition(handle, position)
             .into_iter()
             .next()
-            .map(|r| Range::new(r.module.display_range(r.range)))
+            .map(|r| {
+                // Find the filename corresponding to the module of the definition
+                let target_filename = self
+                    .handles
+                    .iter()
+                    .find(|(_, h)| h.module() == r.module.name())
+                    .map(|(filename, _)| filename.clone())
+                    .unwrap_or_else(|| self.active_filename.clone()); // Fallback to current file
+
+                DefinitionResult {
+                    range: Range::new(r.module.display_range(r.range)),
+                    filename: target_filename,
+                }
+            })
     }
 
     pub fn autocomplete(&self, pos: Position) -> Vec<AutoCompletionItem> {
@@ -523,6 +542,40 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_cross_file_goto_definition() {
+        let mut state = Playground::new(None).unwrap();
+        let mut files = HashMap::new();
+        files.insert(
+            "sandbox.py".to_owned(),
+            "from utils import format_number\n\ndef test(x: int):\n  return format_number(x)"
+                .to_owned(),
+        );
+        files.insert(
+            "utils.py".to_owned(),
+            "def format_number(x: int) -> str:\n    return f\"Number: {x}\"".to_owned(),
+        );
+
+        state.update_sandbox_files(files);
+        state.set_active_file("sandbox.py");
+
+        let result = state.goto_definition(Position::new(4, 10));
+
+        assert!(
+            result.is_some(),
+            "Should find definition for cross-file function call"
+        );
+        let def_result = result.unwrap();
+        assert_eq!(
+            def_result.filename, "utils.py",
+            "Definition should be in utils.py"
+        );
+        assert_eq!(
+            def_result.range.start_line, 1,
+            "Definition should be on line 1"
+        );
+    }
+
     #[test]
     fn test_incremental_update_with_cross_file_errors() {
         let mut state = Playground::new(None).unwrap();

website/src/__tests__/pyrefly_wasm.test.ts

@@ -118,10 +118,10 @@ movie: Movie = {'name': 'Blade Runner',
             expect(definition).toBeDefined();
 
             // expect that location is correct
-            expect(definition.startLineNumber).toBe(13);
-            expect(definition.startColumn).toBe(5);
-            expect(definition.endLineNumber).toBe(13);
-            expect(definition.endColumn).toBe(9);
+            expect(definition.range.startLineNumber).toBe(13);
+            expect(definition.range.startColumn).toBe(5);
+            expect(definition.range.endLineNumber).toBe(13);
+            expect(definition.range.endColumn).toBe(9);
         });
     });
 

website/src/sandbox/Sandbox.tsx

@@ -41,7 +41,7 @@ export interface PyreflyState {
     setActiveFile: (filename: string) => void;
     getErrors: () => ReadonlyArray<PyreflyErrorMessage>;
     autoComplete: (line: number, column: number) => any;
-    gotoDefinition: (line: number, column: number) => any;
+    gotoDefinition: (line: number, column: number) => DefinitionResult | null;
     queryType: (line: number, column: number) => any;
     inlayHint: () => any;
 }
@@ -63,6 +63,17 @@ export async function initializePyreflyWasm(): Promise<any> {
     }
 }
 
+// Types for Pyrefly responses
+export interface DefinitionResult {
+    range: {
+        startLineNumber: number;
+        startColumn: number;
+        endLineNumber: number;
+        endColumn: number;
+    };
+    filename: string;
+}
+
 // This will be used in the component
 let pyreflyWasmInitializedPromise: Promise<any> | null = null;
 
@@ -474,9 +485,16 @@ export default function Sandbox({
         setAutoCompleteFunction(model, (l: number, c: number) =>
             pyreService.autoComplete(l, c)
         );
-        setGetDefFunction(model, (l: number, c: number) =>
-            pyreService.gotoDefinition(l, c)
-        );
+        setGetDefFunction(model, (l: number, c: number) => {
+            const result = pyreService.gotoDefinition(l, c);
+
+            if (result && result.filename !== activeFileName) {
+                setTimeout(() => {
+                    switchToFile(result.filename);
+                }, 100);
+            }
+            return result;
+        });
         setHoverFunctionForMonaco(model, (l: number, c: number) =>
             pyreService.queryType(l, c)
         );

website/src/sandbox/configured-monaco.ts

@@ -15,8 +15,13 @@ type Range = monaco.IRange;
 type Hover = monaco.languages.Hover;
 type InlayHint = monaco.languages.InlayHint;
 
+interface DefinitionResult {
+    range: Range;
+    filename: string;
+}
+
 type AutoCompleteFunction = (line: number, column: number) => CompletionItem[];
-type GetDefFunction = (line: number, column: number) => Range | null;
+type GetDefFunction = (line: number, column: number) => DefinitionResult | null;
 type HoverFunction = (line: number, column: number) => Hover | null;
 type InlayHintFunction = () => InlayHint[];
 
@@ -42,7 +47,7 @@ function setAutoCompleteFunction(
 const defaultGetDefFunctionForMonaco: GetDefFunction = (
     _l: number,
     _c: number
-): Range | null => null;
+): DefinitionResult | null => null;
 const getDefFunctionsForMonaco = new Map<
     monaco.editor.ITextModel,
     GetDefFunction
@@ -85,6 +90,12 @@ function setInlayHintFunctionForMonaco(
     inlayHintFunctionsForMonaco.set(model, f);
 }
 
+function findModelByFilename(filename: string): monaco.editor.ITextModel | null {
+    const models = monaco.editor.getModels();
+    const foundModel = models.find(model => model.uri.path === `/${filename}`);
+    return foundModel || null;
+}
+
 monaco.languages.register({
     id: 'python',
     extensions: ['.py'],
@@ -186,9 +197,18 @@ monaco.languages.registerDefinitionProvider('python', {
             const f =
                 getDefFunctionsForMonaco.get(model) ??
                 defaultGetDefFunctionForMonaco;
-            const range = f(position.lineNumber, position.column);
-            return range != null ? { uri: model.uri, range } : null;
-        } catch (e) {
+            const result = f(position.lineNumber, position.column);
+
+            if (result != null) {
+                // Find the target model based on the filename
+                const targetModel = findModelByFilename(result.filename);
+                const targetUri = targetModel ? targetModel.uri : model.uri;
+
+                return { uri: targetUri, range: result.range };
+            }
+
+            return null;
+            } catch (e) {
             console.error(e);
             return null;
         }