diff --git a/pyrefly/lib/playground.rs b/pyrefly/lib/playground.rs index 759a199ca..6706e2eb5 100644 --- a/pyrefly/lib/playground.rs +++ b/pyrefly/lib/playground.rs @@ -192,6 +192,12 @@ pub struct InlayHint { position: Position, } +#[derive(Serialize)] +pub struct DefinitionResult { + pub range: Range, + pub filename: String, +} + pub struct Playground { state: State, handles: SmallMap, @@ -422,7 +428,7 @@ impl Playground { }) } - pub fn goto_definition(&mut self, pos: Position) -> Option { + pub fn goto_definition(&mut self, pos: Position) -> Option { let handle = self.handles.get(&self.active_filename)?; let transaction = self.state.transaction(); let position = self.to_text_size(&transaction, pos)?; @@ -431,7 +437,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 { @@ -611,6 +630,40 @@ mod tests { } } + #[test] + fn test_cross_file_goto_definition() { + let mut state = Playground::new(None).unwrap(); + let mut files = SmallMap::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(), + "# Utility functions\n# Some comments\n# More comments\n\ndef format_number(x: int) -> str:\n return f\"Number: {x}\"".to_owned(), + ); + + state.update_sandbox_files(files, true); + 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, 5, + "Definition should be on line 5" + ); + } + #[test] fn test_incremental_update_with_cross_file_errors() { let mut state = Playground::new(None).unwrap(); diff --git a/website/src/__tests__/pyrefly_wasm.test.ts b/website/src/__tests__/pyrefly_wasm.test.ts index e49dae2fb..761e64434 100644 --- a/website/src/__tests__/pyrefly_wasm.test.ts +++ b/website/src/__tests__/pyrefly_wasm.test.ts @@ -124,10 +124,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); }); }); diff --git a/website/src/sandbox/Sandbox.tsx b/website/src/sandbox/Sandbox.tsx index 9fc233f61..3b767e37a 100644 --- a/website/src/sandbox/Sandbox.tsx +++ b/website/src/sandbox/Sandbox.tsx @@ -40,7 +40,7 @@ export interface PyreflyState { setActiveFile: (filename: string) => void; getErrors: () => ReadonlyArray; 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; } @@ -62,6 +62,17 @@ export async function initializePyreflyWasm(): Promise { } } +// 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 | null = null; @@ -495,9 +506,25 @@ 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) { + switchToFile(result.filename); + setTimeout(() => { + const editor = editorRef.current; + if (editor) { + editor.setPosition({ + lineNumber: result.range.startLineNumber, + column: result.range.startColumn + }); + editor.revealLineInCenter(result.range.startLineNumber); + } + }, 50); + return null; + } + return result; + }); setHoverFunctionForMonaco(model, (l: number, c: number) => pyreService.queryType(l, c) ); diff --git a/website/src/sandbox/configured-monaco.ts b/website/src/sandbox/configured-monaco.ts index 8c5873a02..65536b69e 100644 --- a/website/src/sandbox/configured-monaco.ts +++ b/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; }