diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 69d97d425..2becc2727 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -1710,11 +1710,51 @@ impl<'a> Transaction<'a> { && definition_range == definition_export.location { references.push(key.range()); + } else if let Key::Definition(id) = key { + // Check if this is a different definition of the same variable name + if module.code_at(id.range()).trim() == module.code_at(definition_range).trim() + && id.range() != definition_range + { + // Only consider it a reference if it's in the same scope level as the original definition + if self.are_in_same_scope(&handle, definition_range, id.range()) { + references.push(id.range()); + } + } } } + Some(references) } + fn are_in_same_scope(&self, handle: &Handle, range1: TextRange, range2: TextRange) -> bool { + let Some(ast) = self.get_ast(handle) else { + return false; + }; + + let scope1 = Self::find_scope_node(&ast, range1.start()); + let scope2 = Self::find_scope_node(&ast, range2.start()); + + scope1 == scope2 + } + + fn find_scope_node(ast: &ModModule, position: TextSize) -> Option { + let covering_nodes = Ast::locate_node(ast, position); + + for node in covering_nodes.iter() { + match node { + AnyNodeRef::StmtFunctionDef(func_def) => { + return Some(func_def.range()); + } + AnyNodeRef::StmtClassDef(class_def) => { + return Some(class_def.range()); + } + _ => continue, + } + } + + None + } + /// Bindings can contain synthetic bindings, which are not meaningful to end users. /// This function helps to filter out such bindings and only leave bindings that eventually /// jumps to a name in the source. diff --git a/pyrefly/lib/test/lsp/local_find_refs.rs b/pyrefly/lib/test/lsp/local_find_refs.rs index 98d995caf..b889d174b 100644 --- a/pyrefly/lib/test/lsp/local_find_refs.rs +++ b/pyrefly/lib/test/lsp/local_find_refs.rs @@ -299,7 +299,6 @@ References: ); } -// todo(kylei): references on renames should find lhs of assignment #[test] fn reassigned_local() { let code = r#" @@ -316,6 +315,8 @@ xy = xy + 1 References: 2 | xy = 5 ^^ +4 | xy = xy + 1 + ^^ 4 | xy = xy + 1 ^^ "# @@ -323,3 +324,56 @@ References: report.trim(), ); } + +#[test] +fn reassigned_local_out_of_scope() { + let code = r#" +xy = 5 +#^ +def increment_xy(): + xy = xy + 1 +"#; + let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report); + assert_eq!( + r#" +# main.py +2 | xy = 5 + ^ +References: +2 | xy = 5 + ^^ +"# + .trim(), + report.trim(), + ); +} + +#[test] +fn reassigned_global() { + let code = r#" +xy = 5 +#^ +def increment_xy(): + global xy + xy = xy + 1 +"#; + let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report); + assert_eq!( + r#" +# main.py +2 | xy = 5 + ^ +References: +2 | xy = 5 + ^^ +5 | global xy + ^^ +6 | xy = xy + 1 + ^^ +6 | xy = xy + 1 + ^^ +"# + .trim(), + report.trim(), + ); +}