From 8e1bb098c5df9920ed7c871a51333beeb8a86b49 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Tue, 27 Oct 2026 20:29:30 +0900 Subject: [PATCH 1/2] test --- .../test/lsp/lsp_interaction/references.rs | 79 ++++++++ .../references_third_party/pyrefly.toml | 1 + .../test_files/references_third_party/user.py | 5 + .../site-packages/third_party_lib.py | 2 + .../site-packages/third_party_lib.pyi | 2 + pyrefly/lib/test/lsp/rename.rs | 181 ++++++++++++++++++ 6 files changed, 270 insertions(+) create mode 100644 pyrefly/lib/test/lsp/lsp_interaction/test_files/references_third_party/pyrefly.toml create mode 100644 pyrefly/lib/test/lsp/lsp_interaction/test_files/references_third_party/user.py create mode 100644 pyrefly/lib/test/lsp/lsp_interaction/test_files/references_third_party/venv/lib/python3.12/site-packages/third_party_lib.py create mode 100644 pyrefly/lib/test/lsp/lsp_interaction/test_files/references_third_party/venv/lib/python3.12/site-packages/third_party_lib.pyi diff --git a/pyrefly/lib/test/lsp/lsp_interaction/references.rs b/pyrefly/lib/test/lsp/lsp_interaction/references.rs index 1cb0122b32..ed4d00dddb 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/references.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/references.rs @@ -683,6 +683,85 @@ fn test_references_for_aliased_import_with_config() { interaction.shutdown().unwrap(); } +#[test] +fn test_references_from_third_party_source_with_stub() { + let root = get_test_files_root(); + let root_path = root.path().join("references_third_party"); + let scope_uri = Uri::from_file_path(root_path.clone()).unwrap(); + let mut interaction = LspInteraction::new_with_args(LspInteractionArgs { + args: LspArgs { + indexing_mode: IndexingMode::LazyBlocking, + ..LspInteractionArgs::default().args + }, + ..Default::default() + }); + interaction.set_root(root_path.clone()); + interaction + .initialize(InitializeSettings { + workspace_folders: Some(vec![("test".to_owned(), scope_uri)]), + ..Default::default() + }) + .unwrap(); + + let library = root_path.join("venv/lib/python3.12/site-packages/third_party_lib.py"); + let user_code = root_path.join("user.py"); + + interaction.client.did_open("user.py"); + interaction + .client + .did_open("venv/lib/python3.12/site-packages/third_party_lib.py"); + + interaction + .client + .references( + "venv/lib/python3.12/site-packages/third_party_lib.py", + 0, + 6, + true, + ) + .expect_response(json!([ + { + "range": {"start":{"line":2,"character":11},"end":{"line":2,"character":16}}, + "uri": Uri::from_file_path(&user_code).unwrap().to_string() + }, + { + "range": {"start":{"line":2,"character":23},"end":{"line":2,"character":28}}, + "uri": Uri::from_file_path(&user_code).unwrap().to_string() + }, + { + "range": {"start":{"line":3,"character":12},"end":{"line":3,"character":17}}, + "uri": Uri::from_file_path(&user_code).unwrap().to_string() + }, + { + "range": {"start":{"line":0,"character":6},"end":{"line":0,"character":11}}, + "uri": Uri::from_file_path(&library).unwrap().to_string() + }, + ])) + .unwrap(); + + interaction + .client + .references( + "venv/lib/python3.12/site-packages/third_party_lib.py", + 1, + 8, + true, + ) + .expect_response(json!([ + { + "range": {"start":{"line":4,"character":6},"end":{"line":4,"character":13}}, + "uri": Uri::from_file_path(&user_code).unwrap().to_string() + }, + { + "range": {"start":{"line":1,"character":8},"end":{"line":1,"character":15}}, + "uri": Uri::from_file_path(&library).unwrap().to_string() + }, + ])) + .unwrap(); + + interaction.shutdown().unwrap(); +} + #[test] fn test_references_after_file_modification_with_config() { let root = get_test_files_root(); diff --git a/pyrefly/lib/test/lsp/lsp_interaction/test_files/references_third_party/pyrefly.toml b/pyrefly/lib/test/lsp/lsp_interaction/test_files/references_third_party/pyrefly.toml new file mode 100644 index 0000000000..a70ea03c2e --- /dev/null +++ b/pyrefly/lib/test/lsp/lsp_interaction/test_files/references_third_party/pyrefly.toml @@ -0,0 +1 @@ +site_package_path = ["venv/lib/python3.12/site-packages"] diff --git a/pyrefly/lib/test/lsp/lsp_interaction/test_files/references_third_party/user.py b/pyrefly/lib/test/lsp/lsp_interaction/test_files/references_third_party/user.py new file mode 100644 index 0000000000..4fc135c906 --- /dev/null +++ b/pyrefly/lib/test/lsp/lsp_interaction/test_files/references_third_party/user.py @@ -0,0 +1,5 @@ +import third_party_lib as lib + +first: lib.Array = lib.Array() +second: lib.Array +first.reshape() diff --git a/pyrefly/lib/test/lsp/lsp_interaction/test_files/references_third_party/venv/lib/python3.12/site-packages/third_party_lib.py b/pyrefly/lib/test/lsp/lsp_interaction/test_files/references_third_party/venv/lib/python3.12/site-packages/third_party_lib.py new file mode 100644 index 0000000000..f2f71a4dc0 --- /dev/null +++ b/pyrefly/lib/test/lsp/lsp_interaction/test_files/references_third_party/venv/lib/python3.12/site-packages/third_party_lib.py @@ -0,0 +1,2 @@ +class Array: + def reshape(self): ... diff --git a/pyrefly/lib/test/lsp/lsp_interaction/test_files/references_third_party/venv/lib/python3.12/site-packages/third_party_lib.pyi b/pyrefly/lib/test/lsp/lsp_interaction/test_files/references_third_party/venv/lib/python3.12/site-packages/third_party_lib.pyi new file mode 100644 index 0000000000..413a235239 --- /dev/null +++ b/pyrefly/lib/test/lsp/lsp_interaction/test_files/references_third_party/venv/lib/python3.12/site-packages/third_party_lib.pyi @@ -0,0 +1,2 @@ +class Array: + def reshape(self) -> None: ... diff --git a/pyrefly/lib/test/lsp/rename.rs b/pyrefly/lib/test/lsp/rename.rs index 319f91d29b..ea314f6a66 100644 --- a/pyrefly/lib/test/lsp/rename.rs +++ b/pyrefly/lib/test/lsp/rename.rs @@ -5,16 +5,197 @@ * LICENSE file in the root directory of this source tree. */ +use std::fs; +use std::ops::Range; + use itertools::Itertools; use pretty_assertions::assert_eq; use pyrefly_build::handle::Handle; +use pyrefly_python::module::TextRangeWithModule; +use pyrefly_python::module_path::ModulePath; use ruff_text_size::TextSize; +use tempfile::TempDir; +use crate::state::lsp::FindPreference; use crate::state::lsp::ReferenceOptions; use crate::state::state::State; +use crate::test::util::TestEnv; use crate::test::util::code_frame_of_source_at_range; use crate::test::util::get_batched_lsp_operations_report; +/// Applies the reference edits used by rename to a source file, its stub, and a consumer. +fn rename_paired_symbol( + source: &str, + stub: &str, + consumer: &str, + position: usize, +) -> Vec<(String, String)> { + let root = TempDir::new().unwrap(); + let mut env = TestEnv::new(); + for (name, file, code) in [ + ("lib", "lib.py", source), + ("lib", "lib.pyi", stub), + ("consumer", "consumer.py", consumer), + ] { + let path = root.path().join(file); + fs::write(&path, code).unwrap(); + env.add_real_path(name, path); + } + let (state, handle) = env.to_state(); + let stub_handle = handle("lib"); + let source_handle = Handle::new( + stub_handle.module(), + ModulePath::filesystem(root.path().join("lib.py")), + *stub_handle.sys_info(), + ); + let mut transaction = state.transaction(); + let definitions = transaction + .find_definition( + &source_handle, + TextSize::try_from(position).unwrap(), + FindPreference { + resolve_call_dunders: false, + ..Default::default() + }, + ) + .unwrap(); + assert_eq!(definitions.len(), 1); + let definition = definitions.first(); + let references = transaction + .find_global_references_from_definition( + *source_handle.sys_info(), + definition.metadata.clone(), + TextRangeWithModule::new(definition.module.clone(), definition.definition_range), + ReferenceOptions::textual_only(true), + ) + .unwrap(); + let mut result = Vec::new(); + for (module, ranges) in references { + let mut code = module.contents().as_str().to_owned(); + for range in ranges.into_iter().rev() { + code.replace_range(Range::::from(range), "renamed"); + } + result.push(( + module + .path() + .as_path() + .file_name() + .unwrap() + .to_str() + .unwrap() + .to_owned(), + code, + )); + } + result.sort(); + result +} + +#[test] +fn test_rename_source_updates_stub_and_consumers() { + for (source, stub, consumer, name) in [ + ( + "def item() -> int:\n return 1\n", + "def item() -> int: ...\n", + "from lib import item\nitem()\n", + "item", + ), + ( + "class Item:\n pass\n", + "class Item: ...\n", + "from lib import Item\nx: Item = Item()\n", + "Item", + ), + ( + "item = 1\n", + "item: int\n", + "from lib import item\nx = item\n", + "item", + ), + ( + "item: int = 1\n", + "item: int\n", + "from lib import item\nx = item\n", + "item", + ), + ( + "class Item:\n def item(self) -> int:\n return 1\n", + "class Item:\n def item(self) -> int: ...\n", + "from lib import Item\nItem().item()\n", + "item", + ), + ( + "class Item:\n item: int = 1\n", + "class Item:\n item: int\n", + "from lib import Item\nx = Item.item\n", + "item", + ), + ( + "class Outer:\n class Item:\n pass\n", + "class Outer:\n class Item: ...\n", + "from lib import Outer\nx: Outer.Item = Outer.Item()\n", + "Item", + ), + ] { + assert_eq!( + rename_paired_symbol(source, stub, consumer, source.find(name).unwrap()), + vec![ + ("consumer.py".to_owned(), consumer.replace(name, "renamed")), + ("lib.py".to_owned(), source.replace(name, "renamed")), + ("lib.pyi".to_owned(), stub.replace(name, "renamed")), + ], + ); + } +} + +#[test] +fn test_rename_local_symbols_does_not_change_stub_or_consumers() { + for (body, position) in [ + ( + "def wrapper(item: int) -> int:\n return item\n", + "item: int", + ), + ( + "def wrapper() -> int:\n item = 1\n return item\n", + "item =", + ), + ( + "def wrapper() -> int:\n def item() -> int:\n return 1\n return item()\n", + "item()", + ), + ] { + for class_scope in [false, true] { + let (source, stub, consumer) = if class_scope { + ( + format!( + "class Item:\n def item(self) -> int:\n return 1\n{}", + body.replacen("wrapper(", "wrapper(self, ", 1) + .lines() + .map(|line| format!(" {line}\n")) + .collect::() + ), + "class Item:\n def item(self) -> int: ...\n", + "from lib import Item\nItem().item()\n", + ) + } else { + ( + format!("def item() -> int:\n return 1\n{body}"), + "def item() -> int: ...\n", + "from lib import item\nitem()\n", + ) + }; + let wrapper_start = source.find("def wrapper").unwrap(); + let start = wrapper_start + source[wrapper_start..].find(position).unwrap(); + let mut expected = source[..start].to_owned(); + expected.push_str(&source[start..].replace("item", "renamed")); + assert_eq!( + rename_paired_symbol(&source, stub, consumer, start), + vec![("lib.py".to_owned(), expected)], + ); + } + } +} + fn get_test_report(state: &State, handle: &Handle, position: TextSize) -> String { let transaction = state.transaction(); // Mirror the reference-collection half of `textDocument/rename`: only ranges whose text From f2120ac584441b7f8667e8c9d4f8030e1aae8b42 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sat, 18 Jul 2026 10:49:15 +0900 Subject: [PATCH 2/2] fix --- pyrefly/lib/lsp/non_wasm/server.rs | 5 +- pyrefly/lib/state/lsp.rs | 169 ++++++++++++++++++++-------- pyrefly/lib/state/lsp_attributes.rs | 51 +++++---- 3 files changed, 153 insertions(+), 72 deletions(-) diff --git a/pyrefly/lib/lsp/non_wasm/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs index 1e9eee5457..826cb697d0 100644 --- a/pyrefly/lib/lsp/non_wasm/server.rs +++ b/pyrefly/lib/lsp/non_wasm/server.rs @@ -5750,7 +5750,10 @@ impl Server { import_behavior: ImportBehavior::StopAtRenamedImports, ..Default::default() }, - options: ReferenceOptions::all(params.context.include_declaration), + options: ReferenceOptions { + include_interface_declaration: false, + ..ReferenceOptions::all(params.context.include_declaration) + }, activity_key, }, move |results| { diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 8088d344cc..8d82f0e031 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -321,6 +321,8 @@ impl Default for FindPreference { pub struct ReferenceOptions { /// Include the definition itself in the results. pub include_declaration: bool, + /// Include the corresponding stub declaration when declarations are requested. + pub include_interface_declaration: bool, /// Include call sites that reach the definition implicitly through the constructor /// protocol, e.g. `Foo()` as a reference to `Foo.__init__`. /// @@ -336,6 +338,7 @@ impl ReferenceOptions { pub fn all(include_declaration: bool) -> Self { Self { include_declaration, + include_interface_declaration: true, include_constructor_call_sites: true, } } @@ -345,6 +348,7 @@ impl ReferenceOptions { pub fn textual_only(include_declaration: bool) -> Self { Self { include_declaration, + include_interface_declaration: true, include_constructor_call_sites: false, } } @@ -2063,6 +2067,39 @@ impl<'a> Transaction<'a> { Some((executable_module, def_range, docstring_range)) } + /// Find the definition in the interface used to type-check an executable module. + fn find_corresponding_interface_definition( + &self, + sys_info: SysInfo, + definition: &TextRangeWithModule, + ) -> Option { + if definition.module.path().style() != ModuleStyle::Executable { + return None; + } + let source_handle = Handle::new( + definition.module.name(), + definition.module.path().dupe(), + sys_info, + ); + let interface_handle = self + .import_handle(&source_handle, definition.module.name(), None) + .finding()?; + if interface_handle.path().style() != ModuleStyle::Interface { + return None; + } + let context = AttributeContext::from_module(&definition.module, definition.range)?; + let name = Name::new(definition.module.code_at(definition.range)); + let _ = self.get_exports(&interface_handle); + let interface_module = self.get_module_info(&interface_handle)?; + let ast = self.get_ast(&interface_handle).unwrap_or_else(|| { + Ast::parse(interface_module.contents(), interface_module.source_type()) + .0 + .into() + }); + let (range, _) = definition_from_executable_ast(ast.as_ref(), &context, &name)?; + Some(TextRangeWithModule::new(interface_module, range)) + } + pub fn key_to_export( &self, handle: &Handle, @@ -5259,6 +5296,11 @@ trait RdepTransaction { fn module_info(&self, handle: &Handle) -> Option; fn transitive_rdeps(&self, handle: Handle) -> HashSet; fn run_for_handles(&mut self, handles: &[Handle], require: Require) -> Result<(), Cancelled>; + fn find_corresponding_interface_definition( + &self, + sys_info: SysInfo, + definition: &TextRangeWithModule, + ) -> Option; fn local_references_from_definition( &self, handle: &Handle, @@ -5288,6 +5330,14 @@ impl<'a> RdepTransaction for Transaction<'a> { Ok(()) } + fn find_corresponding_interface_definition( + &self, + sys_info: SysInfo, + definition: &TextRangeWithModule, + ) -> Option { + self.find_corresponding_interface_definition(sys_info, definition) + } + fn local_references_from_definition( &self, handle: &Handle, @@ -5319,6 +5369,15 @@ impl<'a> RdepTransaction for CancellableTransaction<'a> { self.run(handles, require, None) } + fn find_corresponding_interface_definition( + &self, + sys_info: SysInfo, + definition: &TextRangeWithModule, + ) -> Option { + self.as_ref() + .find_corresponding_interface_definition(sys_info, definition) + } + fn local_references_from_definition( &self, handle: &Handle, @@ -5489,57 +5548,75 @@ fn find_global_references_from_definition_impl( definition: TextRangeWithModule, options: ReferenceOptions, ) -> Result)>, Cancelled> { - let candidate_handles = - compute_transitive_rdeps_for_definition_impl(transaction, sys_info, &definition)?; - if definition_kind.symbol_kind() == Some(SymbolKind::Parameter) { - // Keyword argument references require each candidate's AST and bindings to resolve the - // callee and refine the argument back to this parameter. - transaction.run_for_handles(&candidate_handles, Require::Everything)?; - } - let results = process_candidate_handles_with_definition_impl( - transaction, - candidate_handles, - &definition, - |transaction, handle, patched_definition| { - let mut module_refs: Vec<(Module, Vec)> = Vec::new(); + let interface_definition = + transaction.find_corresponding_interface_definition(sys_info, &definition); + let definitions = interface_definition + .into_iter() + .map(|definition| { + ( + definition, + ReferenceOptions { + include_declaration: options.include_declaration + && options.include_interface_declaration, + ..options + }, + ) + }) + .chain(iter::once((definition, options))); + let mut results = Vec::new(); + for (definition, options) in definitions { + let candidate_handles = + compute_transitive_rdeps_for_definition_impl(transaction, sys_info, &definition)?; + if definition_kind.symbol_kind() == Some(SymbolKind::Parameter) { + // Keyword argument references require each candidate's AST and bindings to resolve the + // callee and refine the argument back to this parameter. + transaction.run_for_handles(&candidate_handles, Require::Everything)?; + } + results.extend(process_candidate_handles_with_definition_impl( + transaction, + candidate_handles, + &definition, + |transaction, handle, patched_definition| { + let mut module_refs: Vec<(Module, Vec)> = Vec::new(); - let references = transaction - .local_references_from_definition( - handle, - definition_kind.clone(), - patched_definition.range, - &patched_definition.module, - options, - ) - .unwrap_or_default(); - if !references.is_empty() - && let Some(module_info) = transaction.module_info(handle) - { - module_refs.push((module_info, references)); - } + let references = transaction + .local_references_from_definition( + handle, + definition_kind.clone(), + patched_definition.range, + &patched_definition.module, + options, + ) + .unwrap_or_default(); + if !references.is_empty() + && let Some(module_info) = transaction.module_info(handle) + { + module_refs.push((module_info, references)); + } - let child_implementations = - find_child_implementations_impl(transaction, handle, patched_definition); - if !child_implementations.is_empty() - && let Some(module_info) = transaction.module_info(handle) - { - if let Some((_, ranges)) = module_refs - .iter_mut() - .find(|(m, _)| m.path() == module_info.path()) + let child_implementations = + find_child_implementations_impl(transaction, handle, patched_definition); + if !child_implementations.is_empty() + && let Some(module_info) = transaction.module_info(handle) { - ranges.extend(child_implementations); - } else { - module_refs.push((module_info, child_implementations)); + if let Some((_, ranges)) = module_refs + .iter_mut() + .find(|(m, _)| m.path() == module_info.path()) + { + ranges.extend(child_implementations); + } else { + module_refs.push((module_info, child_implementations)); + } } - } - if module_refs.is_empty() { - None - } else { - Some(module_refs) - } - }, - ); + if module_refs.is_empty() { + None + } else { + Some(module_refs) + } + }, + )); + } let mut global_references: Vec<(Module, Vec)> = Vec::new(); for module_refs in results { diff --git a/pyrefly/lib/state/lsp_attributes.rs b/pyrefly/lib/state/lsp_attributes.rs index b53e2171bc..c13a960ff0 100644 --- a/pyrefly/lib/state/lsp_attributes.rs +++ b/pyrefly/lib/state/lsp_attributes.rs @@ -28,6 +28,7 @@ pub(crate) struct AttributeContext { } impl AttributeContext { + /// Finds a module-level or class-level declaration with this exact name range. pub(crate) fn from_module( module: &Module, target_range: TextRange, @@ -43,36 +44,36 @@ impl AttributeContext { target_range: TextRange, ) -> Option { for stmt in body { - match stmt { - Stmt::ClassDef(class_def) => { + let kind = match stmt { + Stmt::ClassDef(class_def) if class_def.name.range == target_range => { + AttributeKind::Class + } + Stmt::ClassDef(class_def) if class_def.range().contains_range(target_range) => { parents.push(class_def.name.id.clone()); - if let Some(ctx) = Self::from_body(&class_def.body, parents, target_range) { - return Some(ctx); - } - parents.pop(); - if class_def.range().contains_range(target_range) { - return Some(AttributeContext { - parent_classes: parents.clone(), - kind: AttributeKind::Class, - }); - } + return Self::from_body(&class_def.body, parents, target_range); } - Stmt::FunctionDef(_) if stmt.range().contains_range(target_range) => { - return Some(AttributeContext { - parent_classes: parents.clone(), - kind: AttributeKind::Function, - }); + Stmt::FunctionDef(function) if function.name.range == target_range => { + AttributeKind::Function } - Stmt::Assign(_) | Stmt::AnnAssign(_) - if stmt.range().contains_range(target_range) => + Stmt::Assign(assign) + if assign.targets.iter().any( + |target| matches!(target, Expr::Name(name) if name.range() == target_range), + ) => { - return Some(AttributeContext { - parent_classes: parents.clone(), - kind: AttributeKind::Assignment, - }); + AttributeKind::Assignment } - _ => {} - } + Stmt::AnnAssign(assign) + if assign.target.range() == target_range + && matches!(assign.target.as_ref(), Expr::Name(_)) => + { + AttributeKind::Assignment + } + _ => continue, + }; + return Some(AttributeContext { + parent_classes: parents.clone(), + kind, + }); } None }