From a0b6ed7b2594a39fcc29d92f801e88b5afc14774 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 9 Nov 2026 00:51:12 +0900 Subject: [PATCH 01/19] Test hover types for __all__ entries --- pyrefly/lib/test/lsp/hover.rs | 40 ++++++++++++++++++++++++++++++++--- 1 file changed, 37 insertions(+), 3 deletions(-) diff --git a/pyrefly/lib/test/lsp/hover.rs b/pyrefly/lib/test/lsp/hover.rs index 78b863e9e4..1a3ad1e79d 100644 --- a/pyrefly/lib/test/lsp/hover.rs +++ b/pyrefly/lib/test/lsp/hover.rs @@ -42,8 +42,7 @@ fn get_test_report(state: &State, handle: &Handle, position: TextSize) -> String } } -// BUG: Hover shows string literal types for __all__ and __slots__ entries. -// It should show (variable) foo: int and (attribute) bar: str (issue #4344). +// BUG: Hover shows a string literal type for the __slots__ entry (issue #4344). #[test] fn hover_symbol_literals() { let code = r#" @@ -65,7 +64,7 @@ class C: 3 | __all__ = ["foo"] ^ ```python -(variable) foo: Literal['foo'] +(variable) foo: int ``` 7 | __slots__ = ("bar",) @@ -79,6 +78,41 @@ Literal['bar'] ); } +#[test] +fn hover_dunder_all_reexport_entries() { + let code = r#" +from lib import foo +__all__ = ["foo"] +# ^ +__all__.append("foo") +# ^ +__all__.extend(["foo"]) +# ^ +__all__.remove("foo") +# ^ +"#; + get_batched_lsp_operations_report( + &[("main", code), ("lib", "foo: int = 1\n")], + |state, handle, position| { + let report = get_test_report(state, handle, position); + assert!(report.contains("(variable) foo: int"), "got: {report}"); + report + }, + ); +} + +#[test] +fn hover_string_matching_dunder_all_entry() { + let code = r#" +foo: int = 1 +value = "foo" +# ^ +__all__ = ["foo"] +"#; + let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report); + assert!(report.contains("Literal['foo']"), "got: {report}"); +} + fn get_test_report_at_verbosity( state: &State, handle: &Handle, From a85faf4779354bfe7a1138ad93d7e492667d5d28 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 9 Nov 2026 00:54:33 +0900 Subject: [PATCH 02/19] Show referenced symbol types when hovering over __all__ entries Use the existing definition lookup to read the exported symbol type. Keep string literal types for strings outside __all__. --- pyrefly/lib/lsp/wasm/hover.rs | 18 ++++++++++++++++-- pyrefly/lib/test/lsp/hover.rs | 22 +++++++++++++++++++++- 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/pyrefly/lib/lsp/wasm/hover.rs b/pyrefly/lib/lsp/wasm/hover.rs index 67d8efd3d9..59f7126ad7 100644 --- a/pyrefly/lib/lsp/wasm/hover.rs +++ b/pyrefly/lib/lsp/wasm/hover.rs @@ -1606,7 +1606,8 @@ pub fn get_hover_with_verbosity( position, ); - let Some(type_) = resolve_hovered_type(transaction, handle, ast.as_deref(), position) else { + let Some(mut type_) = resolve_hovered_type(transaction, handle, ast.as_deref(), position) + else { return keyword_hover; }; @@ -1631,7 +1632,6 @@ pub fn get_hover_with_verbosity( }) }); - let fallback_name_from_type = fallback_hover_name_from_type(&type_); let keyword_argument_identifier = keyword_argument_identifier(transaction, handle, position); let definition = transaction .find_definition( @@ -1655,6 +1655,20 @@ pub fn get_hover_with_verbosity( item.module.code_at(item.definition_range) == identifier.id.as_str() }) }); + // An `__all__` entry refers to a symbol, so hover uses that symbol's type. + if let Some(definition) = &definition + && transaction + .get_exports_data(handle) + .dunder_all_name_at(position) + .is_some() + && let Some(definition_type) = transaction.get_type_at( + &handle.for_module(definition.module.name(), definition.module.path().dupe()), + definition.definition_range.start(), + ) + { + type_ = definition_type; + } + let fallback_name_from_type = fallback_hover_name_from_type(&type_); let interface_definition = definition.as_ref().and_then(|executable| { if keyword_argument_identifier.is_some() { return None; diff --git a/pyrefly/lib/test/lsp/hover.rs b/pyrefly/lib/test/lsp/hover.rs index 1a3ad1e79d..4353b35f04 100644 --- a/pyrefly/lib/test/lsp/hover.rs +++ b/pyrefly/lib/test/lsp/hover.rs @@ -57,7 +57,15 @@ class C: def __init__(self) -> None: self.bar: str = "" "#; - let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report); + let report = get_batched_lsp_operations_report(&[("main", code)], |state, handle, position| { + match get_hover(&state.transaction(), handle, position, false) { + Some(Hover { + contents: Contents::MarkupContent(markup), + .. + }) => markup.value, + _ => "None".to_owned(), + } + }); assert_eq!( r#" # main.py @@ -101,6 +109,17 @@ __all__.remove("foo") ); } +#[test] +fn hover_dunder_all_class_entry() { + let code = r#" +class C: ... +__all__ = ["C"] +# ^ +"#; + let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report); + assert!(report.contains("(class) C: type[C]"), "got: {report}"); +} + #[test] fn hover_string_matching_dunder_all_entry() { let code = r#" @@ -2643,6 +2662,7 @@ def exported() -> int: let position = extract_cursors_for_test(lib_code)[0]; let report = get_test_report(&state, &lib_handle, position); + assert!(report.contains("-> int"), "got: {report}"); assert!( report.contains("Public stub documentation."), "got: {report}" From 6e33a9c4ef47f1994594a4faa9076e75fe6ad31a Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Thu, 12 Nov 2026 04:29:54 +0900 Subject: [PATCH 03/19] Use solved export types for __all__ hover MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Read the current checked module’s final export answer so buffers, stubs, and unopened dependencies supply the correct type. --- pyrefly/lib/lsp/wasm/hover.rs | 16 +---- pyrefly/lib/state/lsp.rs | 7 +++ pyrefly/lib/test/lsp/lsp_interaction/hover.rs | 63 +++++++++++++++++++ 3 files changed, 71 insertions(+), 15 deletions(-) diff --git a/pyrefly/lib/lsp/wasm/hover.rs b/pyrefly/lib/lsp/wasm/hover.rs index 59f7126ad7..f2d9f9c084 100644 --- a/pyrefly/lib/lsp/wasm/hover.rs +++ b/pyrefly/lib/lsp/wasm/hover.rs @@ -1606,8 +1606,7 @@ pub fn get_hover_with_verbosity( position, ); - let Some(mut type_) = resolve_hovered_type(transaction, handle, ast.as_deref(), position) - else { + let Some(type_) = resolve_hovered_type(transaction, handle, ast.as_deref(), position) else { return keyword_hover; }; @@ -1655,19 +1654,6 @@ pub fn get_hover_with_verbosity( item.module.code_at(item.definition_range) == identifier.id.as_str() }) }); - // An `__all__` entry refers to a symbol, so hover uses that symbol's type. - if let Some(definition) = &definition - && transaction - .get_exports_data(handle) - .dunder_all_name_at(position) - .is_some() - && let Some(definition_type) = transaction.get_type_at( - &handle.for_module(definition.module.name(), definition.module.path().dupe()), - definition.definition_range.start(), - ) - { - type_ = definition_type; - } let fallback_name_from_type = fallback_hover_name_from_type(&type_); let interface_definition = definition.as_ref().and_then(|executable| { if keyword_argument_identifier.is_some() { diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index ca5ff3723a..082b87324c 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -86,6 +86,7 @@ use crate::alt::attr::AttrInfo; use crate::binding::binding::Binding; use crate::binding::binding::Key; use crate::binding::binding::KeyAnnotation; +use crate::binding::binding::KeyExport; use crate::binding::metadata::BindingsMetadata; use crate::config::error_kind::ErrorKind; use crate::error::error::ErrorQuickFix; @@ -1489,6 +1490,12 @@ impl<'a> Transaction<'a> { position: TextSize, coerce_callees: bool, ) -> Option { + if let Some((_, name)) = self.get_exports_data(handle).dunder_all_name_at(position) { + return self + .get_solutions(handle)? + .get_hashed_opt(Hashed::new(&KeyExport(name))) + .cloned(); + } let Some(IdentifierWithContext { identifier, context, diff --git a/pyrefly/lib/test/lsp/lsp_interaction/hover.rs b/pyrefly/lib/test/lsp/lsp_interaction/hover.rs index 8db4ad9d6a..b8c6bdd448 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/hover.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/hover.rs @@ -5,13 +5,76 @@ * LICENSE file in the root directory of this source tree. */ +use std::fs; + use lsp_types::Uri; use pyrefly_lsp_test::object_model::InitializeSettings; use pyrefly_lsp_test::object_model::LspInteraction; use serde_json::json; +use tempfile::TempDir; use crate::test::lsp::lsp_interaction::util::get_test_files_root; +#[test] +fn hover_dunder_all_uses_open_buffer_export() { + let root = TempDir::new().unwrap(); + fs::write(root.path().join("main.pyi"), "value: bytes\n").unwrap(); + fs::write( + root.path().join("main.py"), + "value = 1\n__all__ = ['value']\n", + ) + .unwrap(); + let mut interaction = LspInteraction::new(); + interaction.set_root(root.path().to_path_buf()); + interaction.initialize(Default::default()).unwrap(); + interaction.client.did_open("main.py"); + interaction.client.did_change( + "main.py", + "\nvalue = 1\nvalue = 'final'\n__all__ = ['value']\n", + ); + interaction + .client + .hover("main.py", 3, 13) + .expect_hover_response_with_markup(|value| { + value.is_some_and(|text| text.contains("(variable) value: str")) + }) + .unwrap(); + interaction.shutdown().unwrap(); +} + +#[test] +fn hover_dunder_all_uses_unopened_dependency_export() { + for with_stub in [false, true] { + let root = TempDir::new().unwrap(); + fs::write(root.path().join("dep.py"), "value: int = 1\n").unwrap(); + if with_stub { + fs::write(root.path().join("dep.pyi"), "value: str\n").unwrap(); + } + fs::write( + root.path().join("main.py"), + "from dep import value as public\n__all__ = ['public']\n", + ) + .unwrap(); + let mut interaction = LspInteraction::new(); + interaction.set_root(root.path().to_path_buf()); + interaction.initialize(Default::default()).unwrap(); + interaction.client.did_open("main.py"); + let expected = if with_stub { + "value: str" + } else { + "value: int" + }; + interaction + .client + .hover("main.py", 1, 13) + .expect_hover_response_with_markup(|value| { + value.is_some_and(|text| text.contains(expected)) + }) + .unwrap(); + interaction.shutdown().unwrap(); + } +} + #[test] fn test_hover_basic() { let root = get_test_files_root(); From 3f68f536064df49314f21d8ce01f7d74c614b18f Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sat, 14 Nov 2026 12:43:27 +0900 Subject: [PATCH 04/19] Keep export lookup off identifier type queries Resolve __all__ entries only after identifier lookup fails. Cover go-to-type-definition and remove unrelated hover changes. --- pyrefly/lib/lsp/wasm/hover.rs | 2 +- pyrefly/lib/state/lsp.rs | 12 ++++++------ pyrefly/lib/test/lsp/hover.rs | 1 - pyrefly/lib/test/lsp/type_definition.rs | 23 +++++++++++++++++++++++ 4 files changed, 30 insertions(+), 8 deletions(-) diff --git a/pyrefly/lib/lsp/wasm/hover.rs b/pyrefly/lib/lsp/wasm/hover.rs index f2d9f9c084..67d8efd3d9 100644 --- a/pyrefly/lib/lsp/wasm/hover.rs +++ b/pyrefly/lib/lsp/wasm/hover.rs @@ -1631,6 +1631,7 @@ pub fn get_hover_with_verbosity( }) }); + let fallback_name_from_type = fallback_hover_name_from_type(&type_); let keyword_argument_identifier = keyword_argument_identifier(transaction, handle, position); let definition = transaction .find_definition( @@ -1654,7 +1655,6 @@ pub fn get_hover_with_verbosity( item.module.code_at(item.definition_range) == identifier.id.as_str() }) }); - let fallback_name_from_type = fallback_hover_name_from_type(&type_); let interface_definition = definition.as_ref().and_then(|executable| { if keyword_argument_identifier.is_some() { return None; diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 082b87324c..eb8e8cb67c 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -1490,17 +1490,17 @@ impl<'a> Transaction<'a> { position: TextSize, coerce_callees: bool, ) -> Option { - if let Some((_, name)) = self.get_exports_data(handle).dunder_all_name_at(position) { - return self - .get_solutions(handle)? - .get_hashed_opt(Hashed::new(&KeyExport(name))) - .cloned(); - } let Some(IdentifierWithContext { identifier, context, }) = self.identifier_at(handle, position) else { + if let Some((_, name)) = self.get_exports_data(handle).dunder_all_name_at(position) { + return self + .get_solutions(handle)? + .get_hashed_opt(Hashed::new(&KeyExport(name))) + .cloned(); + } return self .type_from_match_wildcard_at_impl(handle, position) .or_else(|| self.type_from_expression_at_impl(handle, position, false)); diff --git a/pyrefly/lib/test/lsp/hover.rs b/pyrefly/lib/test/lsp/hover.rs index 4353b35f04..b2f9651d8f 100644 --- a/pyrefly/lib/test/lsp/hover.rs +++ b/pyrefly/lib/test/lsp/hover.rs @@ -2662,7 +2662,6 @@ def exported() -> int: let position = extract_cursors_for_test(lib_code)[0]; let report = get_test_report(&state, &lib_handle, position); - assert!(report.contains("-> int"), "got: {report}"); assert!( report.contains("Public stub documentation."), "got: {report}" diff --git a/pyrefly/lib/test/lsp/type_definition.rs b/pyrefly/lib/test/lsp/type_definition.rs index 03567397f1..7cdcb862e8 100644 --- a/pyrefly/lib/test/lsp/type_definition.rs +++ b/pyrefly/lib/test/lsp/type_definition.rs @@ -37,6 +37,29 @@ fn get_test_report(state: &State, handle: &Handle, position: TextSize) -> String } } +#[test] +fn dunder_all_entry_uses_exported_type_definition() { + let code = r#" +class Exported: ... +value: Exported = Exported() +__all__ = ["value"] +# ^ +"#; + let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report); + assert_eq!( + r#" +# main.py +4 | __all__ = ["value"] + ^ +Type Definition Result: +2 | class Exported: ... + ^^^^^^^^ +"# + .trim(), + report.trim(), + ); +} + #[test] fn function_goes_to_its_own_def() { let code = r#" From a1af502189b7d2ae25a107ae3efeb58825f6d888 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 9 Nov 2026 01:01:48 +0900 Subject: [PATCH 05/19] Test references in __all__ string entries --- pyrefly/lib/test/lsp/local_find_refs.rs | 95 ++++++++++++++++++++++++- 1 file changed, 93 insertions(+), 2 deletions(-) diff --git a/pyrefly/lib/test/lsp/local_find_refs.rs b/pyrefly/lib/test/lsp/local_find_refs.rs index d1a90d01e7..b9b0b5a1fc 100644 --- a/pyrefly/lib/test/lsp/local_find_refs.rs +++ b/pyrefly/lib/test/lsp/local_find_refs.rs @@ -101,8 +101,6 @@ References: ); } -// BUG: References omit __all__ entries. -// Every cursor should find the definition and both string entries (issue #4344). #[test] fn dunder_all_entries_are_symbol_references() { let code = r#" @@ -124,24 +122,117 @@ __all__.remove("foo") References: 2 | foo = 1 ^^^ +4 | __all__ = ["foo"] + ^^^ +6 | __all__.remove("foo") + ^^^ 4 | __all__ = ["foo"] ^ References: 2 | foo = 1 ^^^ +4 | __all__ = ["foo"] + ^^^ +6 | __all__.remove("foo") + ^^^ 6 | __all__.remove("foo") ^ References: 2 | foo = 1 ^^^ +4 | __all__ = ["foo"] + ^^^ +6 | __all__.remove("foo") + ^^^ "# .trim(), report.trim(), ); } +#[test] +fn dunder_all_references_use_string_contents() { + let code = r#" +foo = 1 +# ^ +__all__ = ["foo", r"foo", """foo"""] +__all__.append("foo") +__all__.extend(["foo"]) +__all__.remove("foo") +other = "foo" +def f(): + foo = 2 + return foo +"#; + get_batched_lsp_operations_report(&[("main", code)], |state, handle, position| { + let transaction = state.transaction(); + let module = transaction.get_module_info(handle).unwrap(); + for include_declaration in [false, true] { + let references = transaction.find_local_references( + handle, + position, + ReferenceOptions::all(include_declaration), + ); + let expected_lines = if include_declaration { + vec![1, 3, 3, 3, 4, 5, 6] + } else { + vec![3, 3, 3, 4, 5, 6] + }; + assert_eq!( + references + .iter() + .map(|range| module.to_lsp_range(*range).start.line) + .collect::>(), + expected_lines, + ); + for range in references { + assert_eq!(module.code_at(range), "foo"); + } + } + String::new() + }); +} + +#[test] +fn dunder_all_references_respect_import_aliases() { + let code = r#" +from lib import foo as alias +# ^ +from other import foo +# ^ +__all__ = ["alias", "foo"] +alias +foo +text = "foo" +"#; + get_batched_lsp_operations_report( + &[("main", code), ("lib", "foo = 1\n"), ("other", "foo = 2\n")], + |state, handle, position| { + let transaction = state.transaction(); + let module = transaction.get_module_info(handle).unwrap(); + let references = + transaction.find_local_references(handle, position, ReferenceOptions::all(true)); + let lines = references + .iter() + .map(|range| module.to_lsp_range(*range).start.line) + .collect::>(); + let expected_name = if module.to_lsp_position(position).line == 1 { + assert_eq!(lines, vec![1, 5, 6]); + "alias" + } else { + assert_eq!(lines, vec![3, 5, 7]); + "foo" + }; + for range in references { + assert_eq!(module.code_at(range), expected_name); + } + String::new() + }, + ); +} + // BUG: A __slots__ entry is disconnected from references to its attribute. // Both cursors should find the slots entry, assignment, and read of self.foo (issue #4344). #[test] From a8c59bc3aeef258ea974f7032deeee03d3f87535 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 9 Nov 2026 01:09:07 +0900 Subject: [PATCH 06/19] Include __all__ strings in reference results Resolve each entry to its definition before adding its string content range. Respect import aliases so references to independent names remain separate. --- pyrefly/lib/export/exports.rs | 47 +++++++++++++----- pyrefly/lib/state/lsp.rs | 42 ++++++++++++++++- pyrefly/lib/test/lsp/local_find_refs.rs | 63 +++++++++++++++++++++++++ 3 files changed, 138 insertions(+), 14 deletions(-) diff --git a/pyrefly/lib/export/exports.rs b/pyrefly/lib/export/exports.rs index 2beb837ff7..58578f1785 100644 --- a/pyrefly/lib/export/exports.rs +++ b/pyrefly/lib/export/exports.rs @@ -375,21 +375,44 @@ impl Exports { /// If `position` is inside a user-specified `__all__` string entry, return its range and name. pub fn dunder_all_name_at(&self, position: TextSize) -> Option<(TextRange, Name)> { + self.dunder_all_names()?.find_map(|(range, name)| { + range + .contains_inclusive(position) + .then(|| (range, name.clone())) + }) + } + + /// Return the ranges and names of string entries in a user-specified `__all__`. + pub fn dunder_all_names(&self) -> Option> { if self.definitions.dunder_all.kind != DunderAllKind::Specified { return None; } - self.definitions - .dunder_all - .entries - .iter() - .find_map(|entry| match entry { - DunderAllEntry::Name(range, name) | DunderAllEntry::Remove(range, name) - if range.contains_inclusive(position) => - { - Some((*range, name.clone())) - } - _ => None, - }) + Some( + self.definitions + .dunder_all + .entries + .iter() + .filter_map(|entry| match entry { + DunderAllEntry::Name(range, name) | DunderAllEntry::Remove(range, name) => { + Some((*range, name)) + } + DunderAllEntry::Module(_, _) => None, + }), + ) + } + + /// Return an import's local definition range and whether it has an explicit alias. + pub fn import_definition(&self, name: &Name) -> Option<(TextRange, bool)> { + let definition = self.definitions.definitions.get(name)?; + match definition.style { + DefinitionStyle::ImportAs(_, _) | DefinitionStyle::ImportAsEq(_) => { + Some((definition.range, true)) + } + DefinitionStyle::Import(_) | DefinitionStyle::ImportModule(_) => { + Some((definition.range, false)) + } + _ => None, + } } /// Whether `name` is defined in this module as a `Final` whose type is declared rather than diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index eb8e8cb67c..c4b353195a 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -2935,10 +2935,21 @@ impl<'a> Transaction<'a> { if let Some((definition_handle, _, export)) = self.resolve_named_import(handle, module_info.name(), name.clone(), preference) { - let definition_module = self.get_module_info(&definition_handle)?; + let local_import = exports.import_definition(&name).filter(|(_, has_alias)| { + match preference.import_behavior { + ImportBehavior::StopAtEverything => true, + ImportBehavior::StopAtRenamedImports => *has_alias, + ImportBehavior::JumpThroughEverything => false, + } + }); + let (definition_module, definition_range) = if let Some((range, _)) = local_import { + (module_info, range) + } else { + (self.get_module_info(&definition_handle)?, export.location) + }; return Some(FindDefinitionItemWithDocstring { metadata: DefinitionMetadata::VariableOrAttribute(export.symbol_kind), - definition_range: export.location, + definition_range, module: definition_module, docstring_range: export.docstring_range, display_name: Some(name.to_string()), @@ -4924,6 +4935,33 @@ impl<'a> Transaction<'a> { &definition_name, )); } + if let Some(entries) = self.get_exports_data(handle).dunder_all_names() + && let Some(module_info) = self.get_module_info(handle) + { + let ast = self.get_ast_or_parse_module(handle, &module_info); + for (range, _) in entries { + if let Some(definition) = self.find_definition_for_dunder_all_entry( + handle, + range.start(), + FindPreference { + import_behavior: ImportBehavior::StopAtRenamedImports, + ..Default::default() + }, + ) && definition.module.path() == module.path() + && definition.definition_range == definition_range + && let Some(content_range) = Ast::locate_node(&ast, range.start()) + .into_iter() + .find_map(|node| match node { + AnyNodeRef::ExprStringLiteral(literal) if literal.range() == range => { + Some(literal.as_single_part_string()?.content_range()) + } + _ => None, + }) + { + references.push(content_range); + } + } + } references.sort_by_key(|range| range.start()); references.dedup(); Some(references) diff --git a/pyrefly/lib/test/lsp/local_find_refs.rs b/pyrefly/lib/test/lsp/local_find_refs.rs index b9b0b5a1fc..e3f42dccc0 100644 --- a/pyrefly/lib/test/lsp/local_find_refs.rs +++ b/pyrefly/lib/test/lsp/local_find_refs.rs @@ -5,13 +5,19 @@ * LICENSE file in the root directory of this source tree. */ +use std::fs; + use itertools::Itertools; use pretty_assertions::assert_eq; use pyrefly_build::handle::Handle; +use pyrefly_python::module::TextRangeWithModule; 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; @@ -233,6 +239,63 @@ text = "foo" ); } +#[test] +fn dunder_all_global_references_respect_definition_identity() { + let root = TempDir::new().unwrap(); + let mut env = TestEnv::new(); + for (name, code) in [ + ("lib", "foo = 1\n"), + ("exports", "from lib import foo\n__all__ = [\"foo\"]\n"), + ( + "renamed", + "from lib import foo as alias\n__all__ = [\"alias\"]\n", + ), + ("unrelated", "foo = 2\n__all__ = [\"foo\"]\n"), + ] { + let path = root.path().join(format!("{name}.py")); + fs::write(&path, code).unwrap(); + env.add_real_path(name, path); + } + let (state, handle) = env.to_state(); + let handle = handle("lib"); + let mut transaction = state.transaction(); + let definitions = transaction + .find_definition(&handle, TextSize::new(0), FindPreference::default()) + .unwrap(); + let definition = definitions.first(); + let mut references = transaction + .find_global_references_from_definition( + *handle.sys_info(), + definition.metadata.clone(), + TextRangeWithModule::new(definition.module.clone(), definition.definition_range), + ReferenceOptions::all(true), + ) + .unwrap() + .into_iter() + .map(|(module, ranges)| { + ( + module.name().to_string(), + ranges + .into_iter() + .map(|range| module.code_at(range).to_owned()) + .collect::>(), + ) + }) + .collect::>(); + references.sort(); + assert_eq!( + references, + vec![ + ( + "exports".to_owned(), + vec!["foo".to_owned(), "foo".to_owned()] + ), + ("lib".to_owned(), vec!["foo".to_owned()]), + ("renamed".to_owned(), vec!["foo".to_owned()]), + ], + ); +} + // BUG: A __slots__ entry is disconnected from references to its attribute. // Both cursors should find the slots entry, assignment, and read of self.foo (issue #4344). #[test] From 4222b1f7cbc859e226382db6bcb22e1aa37d7194 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Thu, 12 Nov 2026 04:30:05 +0900 Subject: [PATCH 07/19] Resolve __all__ references through checked bindings Preserve local alias identity and retain export occurrences in the normal reference index so unopened modules need no navigation queries. --- pyrefly/lib/alt/answers.rs | 14 +++ pyrefly/lib/binding/bindings.rs | 27 +++++ pyrefly/lib/export/exports.rs | 14 --- pyrefly/lib/state/lsp.rs | 101 +++++++--------- .../test/lsp/lsp_interaction/references.rs | 114 ++++++++++++++++++ 5 files changed, 200 insertions(+), 70 deletions(-) diff --git a/pyrefly/lib/alt/answers.rs b/pyrefly/lib/alt/answers.rs index 888eef11b2..243913f98e 100644 --- a/pyrefly/lib/alt/answers.rs +++ b/pyrefly/lib/alt/answers.rs @@ -74,6 +74,7 @@ use crate::solver::solver::VarRecurser; use crate::state::errors::ModuleRanges; use crate::state::ide::IntermediateDefinition; use crate::state::ide::key_to_intermediate_definition; +use crate::state::lsp::dunder_all_binding; use crate::state::state::ModuleChanges; use crate::table; use crate::table_for_each; @@ -1527,6 +1528,19 @@ impl Answers { // the preallocated solutions table represents a complete result set. if let Some(index) = &self.index { let mut index = index.lock(); + // Retain export occurrences so unopened reverse dependencies need no AST or answers. + for entry in bindings.dunder_all_entries() { + if let Some(key) = dunder_all_binding(bindings, &entry.id) + && let Some(IntermediateDefinition::NamedImport(_, module, name, None)) = + key_to_intermediate_definition(bindings, key) + { + index + .externally_defined_variable_references + .entry((module, name)) + .or_default() + .push(entry.range); + } + } // Index bindings with external definitions. for idx in bindings.keys::() { let key = bindings.idx_to_key(idx); diff --git a/pyrefly/lib/binding/bindings.rs b/pyrefly/lib/binding/bindings.rs index 4121c722cb..f5ee7339f9 100644 --- a/pyrefly/lib/binding/bindings.rs +++ b/pyrefly/lib/binding/bindings.rs @@ -36,6 +36,7 @@ use pyrefly_util::gas::Gas; use pyrefly_util::suggest::Candidate; use pyrefly_util::suggest::Search; use pyrefly_util::visit::Visit; +use ruff_python_ast::AnyNodeRef; use ruff_python_ast::AtomicNodeIndex; use ruff_python_ast::Expr; use ruff_python_ast::ExprAttribute; @@ -433,6 +434,8 @@ pub struct Bindings { annotated_metadata_ranges: Vec, /// Ranges of the types that were read from docstrings, sorted and non-overlapping. docstring_type_ranges: Vec, + /// Each entry records an `__all__` name and its range without string delimiters. + dunder_all_entries: Vec, scope_trace: Option, /// Names `del`eted at module scope; retained outside the LSP `scope_trace`. module_deletes: SmallSet, @@ -694,6 +697,7 @@ impl Bindings { annotated_metadata_ranges: Vec::new(), docstring_type_ranges: Vec::new(), scope_trace: None, + dunder_all_entries: Vec::new(), module_deletes: SmallSet::new(), unset_nodes: UnsetNodes::default(), global_declarations: SmallSet::new(), @@ -951,6 +955,10 @@ impl Bindings { } } + pub fn dunder_all_entries(&self) -> &[Identifier] { + &self.dunder_all_entries + } + /// Within the LSP, check if a key exists. /// It may not exist within `if False:` or `if sys.version == 0:` style code. pub fn is_valid_key(&self, k: &Key) -> bool { @@ -1109,6 +1117,24 @@ impl Bindings { use_shape_comments_and_stubs: bool, ) -> Self { let module_range = x.range; + let dunder_all_entries = exports + .dunder_all_names() + .into_iter() + .flatten() + .filter_map(|(range, name)| { + Ast::locate_node(&x, range.start()) + .into_iter() + .find_map(|node| match node { + AnyNodeRef::ExprStringLiteral(literal) if literal.range() == range => { + Some(Identifier::new( + name.clone(), + literal.as_single_part_string()?.content_range(), + )) + } + _ => None, + }) + }) + .collect(); let pytest_info = PytestBindingInfo::from_module(&x); let errors = ErrorCollector::new(module_info.dupe(), error_style); let shape_comments = if use_shape_comments_and_stubs { @@ -1297,6 +1323,7 @@ impl Bindings { module_ranges, annotated_metadata_ranges: merge_ranges(builder.annotated_metadata_ranges), docstring_type_ranges: merge_ranges(builder.docstring_type_ranges), + dunder_all_entries, scope_trace: if enable_trace { Some(scope_trace) } else { diff --git a/pyrefly/lib/export/exports.rs b/pyrefly/lib/export/exports.rs index 58578f1785..50503012e1 100644 --- a/pyrefly/lib/export/exports.rs +++ b/pyrefly/lib/export/exports.rs @@ -401,20 +401,6 @@ impl Exports { ) } - /// Return an import's local definition range and whether it has an explicit alias. - pub fn import_definition(&self, name: &Name) -> Option<(TextRange, bool)> { - let definition = self.definitions.definitions.get(name)?; - match definition.style { - DefinitionStyle::ImportAs(_, _) | DefinitionStyle::ImportAsEq(_) => { - Some((definition.range, true)) - } - DefinitionStyle::Import(_) | DefinitionStyle::ImportModule(_) => { - Some((definition.range, false)) - } - _ => None, - } - } - /// Whether `name` is defined in this module as a `Final` whose type is declared rather than /// inferred, as `FinalName::has_declared_type` describes. pub fn final_has_declared_type(&self, name: &Name) -> bool { diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index c4b353195a..8931636c27 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -87,6 +87,7 @@ use crate::binding::binding::Binding; use crate::binding::binding::Key; use crate::binding::binding::KeyAnnotation; use crate::binding::binding::KeyExport; +use crate::binding::bindings::Bindings; use crate::binding::metadata::BindingsMetadata; use crate::config::error_kind::ErrorKind; use crate::error::error::ErrorQuickFix; @@ -382,6 +383,12 @@ fn definition_matcher( } } +/// Connect an export name to its final local binding, including local import aliases. +pub(crate) fn dunder_all_binding<'a>(bindings: &'a Bindings, name: &Name) -> Option<&'a Key> { + let export = bindings.key_to_idx_hashed_opt(Hashed::new(&KeyExport(name.clone())))?; + Some(bindings.idx_to_key(bindings.get(export).key_idx())) +} + /// Parses a forward-reference string such as `"list[Foo]"` into the expression the binder /// built from it. Only strings whose source text is exactly their value are parsed: escapes /// or implicit concatenation shift the ranges of the names parsed from a string. @@ -2932,24 +2939,22 @@ impl<'a> Transaction<'a> { let exports = self.get_exports_data(handle); let (_entry_range, name) = exports.dunder_all_name_at(position)?; - if let Some((definition_handle, _, export)) = - self.resolve_named_import(handle, module_info.name(), name.clone(), preference) - { - let local_import = exports.import_definition(&name).filter(|(_, has_alias)| { - match preference.import_behavior { - ImportBehavior::StopAtEverything => true, - ImportBehavior::StopAtRenamedImports => *has_alias, - ImportBehavior::JumpThroughEverything => false, - } - }); - let (definition_module, definition_range) = if let Some((range, _)) = local_import { - (module_info, range) - } else { - (self.get_module_info(&definition_handle)?, export.location) - }; + let definition = match preference.import_behavior { + ImportBehavior::JumpThroughEverything => self + .resolve_named_import(handle, module_info.name(), name.clone(), preference) + .map(|(handle, _, export)| (handle, export)), + // Reference queries start from the checked binding and preserve local aliases. + ImportBehavior::StopAtEverything | ImportBehavior::StopAtRenamedImports => { + let answers = self.get_answers(handle)?; + let key = dunder_all_binding(answers.bindings(), &name)?; + self.find_export_for_key(handle, key, preference).ok()? + } + }; + if let Some((definition_handle, export)) = definition { + let definition_module = self.get_module_info(&definition_handle)?; return Some(FindDefinitionItemWithDocstring { metadata: DefinitionMetadata::VariableOrAttribute(export.symbol_kind), - definition_range, + definition_range: export.location, module: definition_module, docstring_range: export.docstring_range, display_name: Some(name.to_string()), @@ -4935,33 +4940,6 @@ impl<'a> Transaction<'a> { &definition_name, )); } - if let Some(entries) = self.get_exports_data(handle).dunder_all_names() - && let Some(module_info) = self.get_module_info(handle) - { - let ast = self.get_ast_or_parse_module(handle, &module_info); - for (range, _) in entries { - if let Some(definition) = self.find_definition_for_dunder_all_entry( - handle, - range.start(), - FindPreference { - import_behavior: ImportBehavior::StopAtRenamedImports, - ..Default::default() - }, - ) && definition.module.path() == module.path() - && definition.definition_range == definition_range - && let Some(content_range) = Ast::locate_node(&ast, range.start()) - .into_iter() - .find_map(|node| match node { - AnyNodeRef::ExprStringLiteral(literal) if literal.range() == range => { - Some(literal.as_single_part_string()?.content_range()) - } - _ => None, - }) - { - references.push(content_range); - } - } - } references.sort_by_key(|range| range.start()); references.dedup(); Some(references) @@ -5180,19 +5158,18 @@ impl<'a> Transaction<'a> { ) -> Option> { let mut references = Vec::new(); if let Some(mod_module) = self.get_ast(handle) { - let is_valid_use = |x: &ExprName| { - if x.id() == expected_name - && let Some((def_handle, Export { location, .. })) = self - .find_export_for_key( - handle, - &Key::BoundName(ShortIdentifier::expr_name(x)), - FindPreference { - import_behavior: ImportBehavior::StopAtRenamedImports, - prefer_pyi: false, - ..Default::default() - }, - ) - .unwrap_or(None) + let is_valid_key = |key: &Key| { + if let Some((def_handle, Export { location, .. })) = self + .find_export_for_key( + handle, + key, + FindPreference { + import_behavior: ImportBehavior::StopAtRenamedImports, + prefer_pyi: false, + ..Default::default() + }, + ) + .unwrap_or(None) && def_handle.path() == handle.path() && location == definition_range { @@ -5201,6 +5178,18 @@ impl<'a> Transaction<'a> { false } }; + if let Some(answers) = self.get_answers(handle) { + let bindings = answers.bindings(); + references.extend(bindings.dunder_all_entries().iter().filter_map(|entry| { + (entry.id == *expected_name + && dunder_all_binding(bindings, &entry.id).is_some_and(is_valid_key)) + .then_some(entry.range) + })); + } + let is_valid_use = |x: &ExprName| { + x.id() == expected_name + && is_valid_key(&Key::BoundName(ShortIdentifier::expr_name(x))) + }; let module_info = self.get_module_info(handle); // Only strings that mention the name can hold a forward reference to it. let forward_ref = |x: &ExprStringLiteral| { diff --git a/pyrefly/lib/test/lsp/lsp_interaction/references.rs b/pyrefly/lib/test/lsp/lsp_interaction/references.rs index 3e5bbe3db3..4a806f68d1 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/references.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/references.rs @@ -20,6 +20,120 @@ use tempfile::TempDir; use crate::test::lsp::lsp_interaction::util::get_test_files_root; use crate::test::lsp::lsp_interaction::util::wait_for_disk_invalidation; +#[test] +fn references_dunder_all_uses_open_buffer_binding() { + let root = TempDir::new().unwrap(); + fs::write(root.path().join("dep.py"), "foo = 42\n").unwrap(); + fs::write(root.path().join("main.py"), "foo = 1\n__all__ = ['foo']\n").unwrap(); + let mut interaction = LspInteraction::new(); + interaction.set_root(root.path().to_path_buf()); + interaction.initialize(Default::default()).unwrap(); + interaction.client.did_open("main.py"); + interaction.client.did_change( + "main.py", + "from dep import foo\nfoo = 1\n__all__ = ['foo']\nfoo\ndef f():\n foo = 2\n return foo\n", + ); + let uri = Uri::from_file_path(root.path().join("main.py")).unwrap(); + for (line, column) in [(1, 1), (2, 13), (3, 1)] { + interaction + .client + .references("main.py", line, column, true) + .expect_response(json!([ + {"uri": uri, "range": {"start": {"line": 1, "character": 0}, "end": {"line": 1, "character": 3}}}, + {"uri": uri, "range": {"start": {"line": 2, "character": 12}, "end": {"line": 2, "character": 15}}}, + {"uri": uri, "range": {"start": {"line": 3, "character": 0}, "end": {"line": 3, "character": 3}}}, + ])) + .unwrap(); + } + interaction.shutdown().unwrap(); +} + +#[test] +fn references_dunder_all_preserves_alias_with_unopened_stub() { + let root = TempDir::new().unwrap(); + fs::write(root.path().join("dep.py"), "\nfoo: int = 1\n").unwrap(); + fs::write(root.path().join("dep.pyi"), "foo: str\n").unwrap(); + fs::write( + root.path().join("main.py"), + "from dep import foo as alias\nfrom dep import foo\n__all__ = ['alias', 'foo']\nalias\nfoo\n", + ) + .unwrap(); + let mut interaction = LspInteraction::new(); + interaction.set_root(root.path().to_path_buf()); + interaction.initialize(Default::default()).unwrap(); + interaction.client.did_open("main.py"); + let uri = Uri::from_file_path(root.path().join("main.py")).unwrap(); + for (line, column) in [(0, 23), (2, 13), (3, 1)] { + interaction + .client + .references("main.py", line, column, true) + .expect_response(json!([ + {"uri": uri, "range": {"start": {"line": 0, "character": 23}, "end": {"line": 0, "character": 28}}}, + {"uri": uri, "range": {"start": {"line": 2, "character": 12}, "end": {"line": 2, "character": 17}}}, + {"uri": uri, "range": {"start": {"line": 3, "character": 0}, "end": {"line": 3, "character": 5}}}, + ])) + .unwrap(); + } + interaction + .client + .references("main.py", 2, 21, false) + .expect_response(json!([ + {"uri": uri, "range": {"start": {"line": 0, "character": 16}, "end": {"line": 0, "character": 19}}}, + {"uri": uri, "range": {"start": {"line": 1, "character": 16}, "end": {"line": 1, "character": 19}}}, + {"uri": uri, "range": {"start": {"line": 2, "character": 21}, "end": {"line": 2, "character": 24}}}, + {"uri": uri, "range": {"start": {"line": 4, "character": 0}, "end": {"line": 4, "character": 3}}}, + ])) + .unwrap(); + interaction.shutdown().unwrap(); +} + +#[test] +fn references_dunder_all_in_unopened_reexports() { + let root = TempDir::new().unwrap(); + for (file, code) in [ + ("dep.py", "foo = 1\n"), + ("exports.py", "from dep import foo\n__all__ = ['foo']\n"), + ( + "renamed.py", + "from dep import foo as alias\n__all__ = ['alias']\n", + ), + ("other.py", "foo = 2\n__all__ = ['foo']\n"), + ] { + fs::write(root.path().join(file), code).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().to_path_buf()); + interaction + .initialize(InitializeSettings { + workspace_folders: Some(vec![( + "test".to_owned(), + Uri::from_file_path(root.path()).unwrap(), + )]), + configuration: Some(None), + ..Default::default() + }) + .unwrap(); + interaction.client.did_open("dep.py"); + let exports_uri = Uri::from_file_path(root.path().join("exports.py")).unwrap(); + let renamed_uri = Uri::from_file_path(root.path().join("renamed.py")).unwrap(); + interaction + .client + .references("dep.py", 0, 1, false) + .expect_response(json!([ + {"uri": exports_uri, "range": {"start": {"line": 0, "character": 16}, "end": {"line": 0, "character": 19}}}, + {"uri": exports_uri, "range": {"start": {"line": 1, "character": 12}, "end": {"line": 1, "character": 15}}}, + {"uri": renamed_uri, "range": {"start": {"line": 0, "character": 16}, "end": {"line": 0, "character": 19}}}, + ])) + .unwrap(); + interaction.shutdown().unwrap(); +} + // https://github.com/meta/pyrefly/issues/5485 #[test] fn test_references_to_class_member_after_defining_file_shifts() { From e36ec6d0c083a0bda48ad33effdfad4cec16b0cf Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sat, 14 Nov 2026 12:53:32 +0900 Subject: [PATCH 08/19] Extract export occurrences once and share binding resolution Keep string content ranges from export extraction and resolve navigation and references through the checked binding. Preserve interface documentation without changing symbol identity. --- pyrefly/lib/alt/answers.rs | 3 +- pyrefly/lib/binding/bindings.rs | 23 ++--- pyrefly/lib/export/definitions.rs | 91 ++++++++++++++----- pyrefly/lib/export/exports.rs | 51 +++++------ pyrefly/lib/lsp/wasm/hover.rs | 25 +++-- pyrefly/lib/state/lsp.rs | 43 ++++----- pyrefly/lib/stubgen/extract.rs | 4 +- pyrefly/lib/test/lsp/local_find_refs.rs | 63 ------------- .../test/lsp/lsp_interaction/references.rs | 5 + 9 files changed, 144 insertions(+), 164 deletions(-) diff --git a/pyrefly/lib/alt/answers.rs b/pyrefly/lib/alt/answers.rs index 243913f98e..4b9534fa3c 100644 --- a/pyrefly/lib/alt/answers.rs +++ b/pyrefly/lib/alt/answers.rs @@ -74,7 +74,6 @@ use crate::solver::solver::VarRecurser; use crate::state::errors::ModuleRanges; use crate::state::ide::IntermediateDefinition; use crate::state::ide::key_to_intermediate_definition; -use crate::state::lsp::dunder_all_binding; use crate::state::state::ModuleChanges; use crate::table; use crate::table_for_each; @@ -1530,7 +1529,7 @@ impl Answers { let mut index = index.lock(); // Retain export occurrences so unopened reverse dependencies need no AST or answers. for entry in bindings.dunder_all_entries() { - if let Some(key) = dunder_all_binding(bindings, &entry.id) + if let Some(key) = bindings.dunder_all_binding(&entry.id) && let Some(IntermediateDefinition::NamedImport(_, module, name, None)) = key_to_intermediate_definition(bindings, key) { diff --git a/pyrefly/lib/binding/bindings.rs b/pyrefly/lib/binding/bindings.rs index f5ee7339f9..e11472d16f 100644 --- a/pyrefly/lib/binding/bindings.rs +++ b/pyrefly/lib/binding/bindings.rs @@ -36,7 +36,6 @@ use pyrefly_util::gas::Gas; use pyrefly_util::suggest::Candidate; use pyrefly_util::suggest::Search; use pyrefly_util::visit::Visit; -use ruff_python_ast::AnyNodeRef; use ruff_python_ast::AtomicNodeIndex; use ruff_python_ast::Expr; use ruff_python_ast::ExprAttribute; @@ -959,6 +958,12 @@ impl Bindings { &self.dunder_all_entries } + /// Return the final local binding for an export, preserving local import aliases. + pub fn dunder_all_binding(&self, name: &Name) -> Option<&Key> { + let export = self.key_to_idx_hashed_opt(Hashed::new(&KeyExport(name.clone())))?; + Some(self.idx_to_key(self.get(export).key_idx())) + } + /// Within the LSP, check if a key exists. /// It may not exist within `if False:` or `if sys.version == 0:` style code. pub fn is_valid_key(&self, k: &Key) -> bool { @@ -1119,20 +1124,8 @@ impl Bindings { let module_range = x.range; let dunder_all_entries = exports .dunder_all_names() - .into_iter() - .flatten() - .filter_map(|(range, name)| { - Ast::locate_node(&x, range.start()) - .into_iter() - .find_map(|node| match node { - AnyNodeRef::ExprStringLiteral(literal) if literal.range() == range => { - Some(Identifier::new( - name.clone(), - literal.as_single_part_string()?.content_range(), - )) - } - _ => None, - }) + .filter_map(|(_, name, content_range)| { + Some(Identifier::new(name.clone(), content_range?)) }) .collect(); let pytest_info = PytestBindingInfo::from_module(&x); diff --git a/pyrefly/lib/export/definitions.rs b/pyrefly/lib/export/definitions.rs index db4014bab2..ec175ae64c 100644 --- a/pyrefly/lib/export/definitions.rs +++ b/pyrefly/lib/export/definitions.rs @@ -277,10 +277,11 @@ pub enum DunderAllKind { #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum DunderAllEntry { - Name(TextRange, Name), + /// The optional content range excludes the delimiters of a single string literal. + Name(TextRange, Name, Option), Module(TextRange, ModuleName), // We have to have this explicitly, as you might remove something in a Module - Remove(TextRange, Name), + Remove(TextRange, Name, Option), } /// The `__all__` variable contents with tracking of whether it was user-specified. @@ -349,9 +350,11 @@ impl DunderAllEntry { fn as_item(x: &Expr) -> Option { match x { - Expr::StringLiteral(x) => { - Some(DunderAllEntry::Name(x.range, Name::new(x.value.to_str()))) - } + Expr::StringLiteral(x) => Some(DunderAllEntry::Name( + x.range, + Name::new(x.value.to_str()), + x.as_single_part_string().map(|part| part.content_range()), + )), _ => None, } } @@ -505,7 +508,7 @@ impl Definitions { { self.dunder_all .entries - .push(DunderAllEntry::Name(def.range, name.clone())); + .push(DunderAllEntry::Name(def.range, name.clone(), None)); } } } @@ -526,7 +529,7 @@ impl Definitions { if let Some(def) = self.definitions.get(name) { self.dunder_all .entries - .push(DunderAllEntry::Name(def.range, name.clone())) + .push(DunderAllEntry::Name(def.range, name.clone(), None)) } } } @@ -986,13 +989,14 @@ impl DefinitionsBuilder { } }, "remove" => { - if let Some(DunderAllEntry::Name(range, remove)) = + if let Some(DunderAllEntry::Name(range, remove, content_range)) = DunderAllEntry::as_item(&arguments.args[0]) { - self.inner - .dunder_all - .entries - .push(DunderAllEntry::Remove(range, remove)); + self.inner.dunder_all.entries.push(DunderAllEntry::Remove( + range, + remove, + content_range, + )); } } _ => {} @@ -1245,9 +1249,12 @@ mod tests { fn unrange(x: &mut DunderAllEntry) { match x { - DunderAllEntry::Name(range, _) - | DunderAllEntry::Module(range, _) - | DunderAllEntry::Remove(range, _) => { + DunderAllEntry::Name(range, _, content_range) + | DunderAllEntry::Remove(range, _, content_range) => { + *range = TextRange::default(); + *content_range = None; + } + DunderAllEntry::Module(range, _) => { *range = TextRange::default(); } } @@ -1419,6 +1426,42 @@ def bar(x: str) -> str: ... assert!(bar.needs_anywhere); } + #[test] + fn test_all_content_ranges() { + let source = "__all__ = [r'foo', \"\"\"bar\"\"\", 'ba' 'z']\n__all__.remove('foo')\n"; + let defs = Definitions::new( + &Ast::parse(source, PySourceType::Python).0.body, + ModuleName::from_str("main"), + false, + SysInfo::default(), + ); + let entries = defs + .dunder_all + .entries + .iter() + .map(|entry| match entry { + DunderAllEntry::Name(range, name, content_range) + | DunderAllEntry::Remove(range, name, content_range) => ( + &source[*range], + name.as_str(), + content_range.map(|range| &source[range]), + ), + DunderAllEntry::Module(..) => { + unreachable!("The fixture contains only string entries") + } + }) + .collect::>(); + assert_eq!( + entries, + vec![ + ("r'foo'", "foo", Some("foo")), + ("\"\"\"bar\"\"\"", "bar", Some("bar")), + ("'ba' 'z'", "baz", None), + ("'foo'", "foo", Some("foo")), + ], + ); + } + #[test] fn test_all() { let defs = calculate_unranged_definitions_with_defaults( @@ -1442,10 +1485,10 @@ __all__.remove('r') assert_eq!(defs.dunder_all.kind, DunderAllKind::Specified); let loc = TextRange::default(); - let a = &DunderAllEntry::Name(loc, Name::new_static("a")); - let b = &DunderAllEntry::Name(loc, Name::new_static("b")); + let a = &DunderAllEntry::Name(loc, Name::new_static("a"), None); + let b = &DunderAllEntry::Name(loc, Name::new_static("b"), None); let foo = &DunderAllEntry::Module(loc, ModuleName::from_str("foo")); - let r = &DunderAllEntry::Remove(loc, Name::new_static("r")); + let r = &DunderAllEntry::Remove(loc, Name::new_static("r"), None); assert_eq!( defs.dunder_all.entries.map(|x| x), vec![a, b, a, b, foo, a, b, foo, a, r] @@ -1483,8 +1526,8 @@ __all__: list[str] = ["a", "b"] assert_definition_names(&defs, &["a", "b", "__all__"]); assert_eq!(defs.dunder_all.kind, DunderAllKind::Specified); let loc = TextRange::default(); - let a = &DunderAllEntry::Name(loc, Name::new_static("a")); - let b = &DunderAllEntry::Name(loc, Name::new_static("b")); + let a = &DunderAllEntry::Name(loc, Name::new_static("a"), None); + let b = &DunderAllEntry::Name(loc, Name::new_static("b"), None); assert_eq!(defs.dunder_all.entries.map(|x| x), vec![a, b]); } @@ -1521,8 +1564,8 @@ __all__ = ["a"] + ["b"] ); assert_eq!(defs.dunder_all.kind, DunderAllKind::Specified); let loc = TextRange::default(); - let a = &DunderAllEntry::Name(loc, Name::new_static("a")); - let b = &DunderAllEntry::Name(loc, Name::new_static("b")); + let a = &DunderAllEntry::Name(loc, Name::new_static("a"), None); + let b = &DunderAllEntry::Name(loc, Name::new_static("b"), None); assert_eq!(defs.dunder_all.entries.map(|x| x), vec![a, b]); } @@ -1536,7 +1579,7 @@ __all__ = ["a"] + foo.__all__ ); assert_eq!(defs.dunder_all.kind, DunderAllKind::Specified); let loc = TextRange::default(); - let a = &DunderAllEntry::Name(loc, Name::new_static("a")); + let a = &DunderAllEntry::Name(loc, Name::new_static("a"), None); let foo = &DunderAllEntry::Module(loc, ModuleName::from_str("foo")); assert_eq!(defs.dunder_all.entries.map(|x| x), vec![a, foo]); } @@ -1552,7 +1595,7 @@ __all__ = [*foo.__all__, "a"] assert_eq!(defs.dunder_all.kind, DunderAllKind::Specified); let loc = TextRange::default(); let foo = &DunderAllEntry::Module(loc, ModuleName::from_str("foo")); - let a = &DunderAllEntry::Name(loc, Name::new_static("a")); + let a = &DunderAllEntry::Name(loc, Name::new_static("a"), None); assert_eq!(defs.dunder_all.entries.map(|x| x), vec![foo, a]); } diff --git a/pyrefly/lib/export/exports.rs b/pyrefly/lib/export/exports.rs index 50503012e1..590db3b2f5 100644 --- a/pyrefly/lib/export/exports.rs +++ b/pyrefly/lib/export/exports.rs @@ -155,9 +155,9 @@ impl Display for Exports { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { for x in self.definitions.dunder_all.entries.iter() { match x { - DunderAllEntry::Name(_, x) => writeln!(f, "export {x}")?, + DunderAllEntry::Name(_, x, _) => writeln!(f, "export {x}")?, DunderAllEntry::Module(_, x) => writeln!(f, "from {x} import *")?, - DunderAllEntry::Remove(_, x) => writeln!(f, "unexport {x}")?, + DunderAllEntry::Remove(_, x, _) => writeln!(f, "unexport {x}")?, } } Ok(()) @@ -241,7 +241,7 @@ impl Exports { let mut result = SmallSet::new(); for x in &self.definitions.dunder_all.entries { match x { - DunderAllEntry::Name(_, x) => { + DunderAllEntry::Name(_, x, _) => { // An empty entry names nothing importable, and the binding pass // does not export empty names, so exporting one here would // promise an export that never exists. @@ -269,7 +269,7 @@ impl Exports { } } } - DunderAllEntry::Remove(_, x) => { + DunderAllEntry::Remove(_, x, _) => { // This is O(n), but we'd appreciate the determism, and remove is rare in `__all__`. result.shift_remove(x); } @@ -375,30 +375,27 @@ impl Exports { /// If `position` is inside a user-specified `__all__` string entry, return its range and name. pub fn dunder_all_name_at(&self, position: TextSize) -> Option<(TextRange, Name)> { - self.dunder_all_names()?.find_map(|(range, name)| { + self.dunder_all_names().find_map(|(range, name, _)| { range .contains_inclusive(position) .then(|| (range, name.clone())) }) } - /// Return the ranges and names of string entries in a user-specified `__all__`. - pub fn dunder_all_names(&self) -> Option> { - if self.definitions.dunder_all.kind != DunderAllKind::Specified { - return None; - } - Some( - self.definitions - .dunder_all - .entries - .iter() - .filter_map(|entry| match entry { - DunderAllEntry::Name(range, name) | DunderAllEntry::Remove(range, name) => { - Some((*range, name)) - } - DunderAllEntry::Module(_, _) => None, - }), - ) + /// Return expression ranges, names, and content ranges from a user-specified `__all__`. + pub fn dunder_all_names(&self) -> impl Iterator)> { + let entries = if self.definitions.dunder_all.kind == DunderAllKind::Specified { + self.definitions.dunder_all.entries.as_slice() + } else { + &[] + }; + entries.iter().filter_map(|entry| match entry { + DunderAllEntry::Name(range, name, content_range) + | DunderAllEntry::Remove(range, name, content_range) => { + Some((*range, name, *content_range)) + } + DunderAllEntry::Module(_, _) => None, + }) } /// Whether `name` is defined in this module as a `Final` whose type is declared rather than @@ -462,10 +459,10 @@ impl Exports { for entry in &definitions.dunder_all.entries { match entry { - DunderAllEntry::Name(_, name) => { + DunderAllEntry::Name(_, name, _) => { names.insert(name.clone()); } - DunderAllEntry::Remove(_, name) => { + DunderAllEntry::Remove(_, name, _) => { names.shift_remove(name); } DunderAllEntry::Module(..) => {} @@ -485,10 +482,10 @@ impl Exports { for entry in &definitions.dunder_all.entries { match entry { - DunderAllEntry::Name(_, name) => { + DunderAllEntry::Name(_, name, _) => { names.insert(name.clone()); } - DunderAllEntry::Remove(_, name) => { + DunderAllEntry::Remove(_, name, _) => { names.shift_remove(name); } DunderAllEntry::Module(..) => {} @@ -512,7 +509,7 @@ impl Exports { } let mut invalid = Vec::new(); for entry in &self.definitions.dunder_all.entries { - if let DunderAllEntry::Name(range, name) = entry { + if let DunderAllEntry::Name(range, name, _) = entry { // Check if name exists in definitions if self.definitions.definitions.contains_key(name) { continue; diff --git a/pyrefly/lib/lsp/wasm/hover.rs b/pyrefly/lib/lsp/wasm/hover.rs index 67d8efd3d9..2fa7c2c0a5 100644 --- a/pyrefly/lib/lsp/wasm/hover.rs +++ b/pyrefly/lib/lsp/wasm/hover.rs @@ -27,6 +27,7 @@ use pyrefly_python::ignore::Ignore; use pyrefly_python::ignore::Tool; use pyrefly_python::ignore::TypeIgnoreUnknownTagBehavior; use pyrefly_python::module::Module; +use pyrefly_python::module::TextRangeWithModule; use pyrefly_python::short_identifier::ShortIdentifier; use pyrefly_python::symbol_kind::SymbolKind; use pyrefly_types::callable::Callable; @@ -1655,7 +1656,7 @@ pub fn get_hover_with_verbosity( item.module.code_at(item.definition_range) == identifier.id.as_str() }) }); - let interface_definition = definition.as_ref().and_then(|executable| { + let interface_docstring = definition.as_ref().and_then(|executable| { if keyword_argument_identifier.is_some() { return None; } @@ -1677,6 +1678,22 @@ pub fn get_hover_with_verbosity( && executable.metadata.symbol_kind() == interface.metadata.symbol_kind() && executable.display_name == interface.display_name }) + .and_then(|interface| { + interface + .docstring_range + .map(|range| Docstring(range, interface.module)) + }) + .or_else(|| { + let (interface, docstring_range) = transaction + .find_corresponding_interface_definition( + *handle.sys_info(), + &TextRangeWithModule::new( + executable.module.dupe(), + executable.definition_range, + ), + )?; + Some(Docstring(docstring_range?, interface.module)) + }) }); let hover_constructor_dunder = definition .as_ref() @@ -1700,11 +1717,7 @@ pub fn get_hover_with_verbosity( } else { (None, fallback_name_from_type, None) }; - let docstring = interface_definition - .and_then(|item| { - item.docstring_range - .map(|range| Docstring(range, item.module)) - }) + let docstring = interface_docstring .and_then(nonempty_docstring) .or_else(|| fallback_docstring.and_then(nonempty_docstring)) .or_else(|| docstring_for_class_object_type(transaction, handle, &type_)); diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 8931636c27..7439c6b0eb 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -87,7 +87,6 @@ use crate::binding::binding::Binding; use crate::binding::binding::Key; use crate::binding::binding::KeyAnnotation; use crate::binding::binding::KeyExport; -use crate::binding::bindings::Bindings; use crate::binding::metadata::BindingsMetadata; use crate::config::error_kind::ErrorKind; use crate::error::error::ErrorQuickFix; @@ -383,12 +382,6 @@ fn definition_matcher( } } -/// Connect an export name to its final local binding, including local import aliases. -pub(crate) fn dunder_all_binding<'a>(bindings: &'a Bindings, name: &Name) -> Option<&'a Key> { - let export = bindings.key_to_idx_hashed_opt(Hashed::new(&KeyExport(name.clone())))?; - Some(bindings.idx_to_key(bindings.get(export).key_idx())) -} - /// Parses a forward-reference string such as `"list[Foo]"` into the expression the binder /// built from it. Only strings whose source text is exactly their value are parsed: escapes /// or implicit concatenation shift the ranges of the names parsed from a string. @@ -2156,12 +2149,12 @@ impl<'a> Transaction<'a> { /// a different file, because the interface then describes that other module. Resolution can /// also find no executable and fall back to the interface, for example for a `-stubs` package; /// the interface is still used then. Open files have in-memory module paths, so the file - /// paths are compared rather than the module paths. - fn find_corresponding_interface_definition( + /// paths are compared rather than the module paths. The result includes the docstring range. + pub(crate) fn find_corresponding_interface_definition( &self, sys_info: SysInfo, definition: &TextRangeWithModule, - ) -> Option { + ) -> Option<(TextRangeWithModule, Option)> { if definition.module.path().style() != ModuleStyle::Executable { return None; } @@ -2194,9 +2187,12 @@ impl<'a> Transaction<'a> { } let context = AttributeContext::from_module(&definition.module, definition.range)?; let name = Name::new(definition.module.code_at(definition.range)); - let (interface_module, range, _) = + let (interface_module, range, docstring_range) = self.find_declaration_in_module(&interface_handle, &context, &name)?; - Some(TextRangeWithModule::new(interface_module, range)) + Some(( + TextRangeWithModule::new(interface_module, range), + docstring_range, + )) } /// Find the declaration of `name` in the scope that `context` describes, within the module @@ -2939,18 +2935,11 @@ impl<'a> Transaction<'a> { let exports = self.get_exports_data(handle); let (_entry_range, name) = exports.dunder_all_name_at(position)?; - let definition = match preference.import_behavior { - ImportBehavior::JumpThroughEverything => self - .resolve_named_import(handle, module_info.name(), name.clone(), preference) - .map(|(handle, _, export)| (handle, export)), - // Reference queries start from the checked binding and preserve local aliases. - ImportBehavior::StopAtEverything | ImportBehavior::StopAtRenamedImports => { - let answers = self.get_answers(handle)?; - let key = dunder_all_binding(answers.bindings(), &name)?; + let answers = self.get_answers(handle)?; + if let Some(key) = answers.bindings().dunder_all_binding(&name) + && let Some((definition_handle, export)) = self.find_export_for_key(handle, key, preference).ok()? - } - }; - if let Some((definition_handle, export)) = definition { + { let definition_module = self.get_module_info(&definition_handle)?; return Some(FindDefinitionItemWithDocstring { metadata: DefinitionMetadata::VariableOrAttribute(export.symbol_kind), @@ -4681,7 +4670,7 @@ impl<'a> Transaction<'a> { // Also block rename if the stub that consumers are checked against cannot be edited, // because renaming the definition without that stub would break its consumers. - if let Some(interface) = self.find_corresponding_interface_definition( + if let Some((interface, _)) = self.find_corresponding_interface_definition( *handle.sys_info(), &TextRangeWithModule::new(module, definition_range), ) && (interface.module.path().is_bundled() @@ -5182,7 +5171,9 @@ impl<'a> Transaction<'a> { let bindings = answers.bindings(); references.extend(bindings.dunder_all_entries().iter().filter_map(|entry| { (entry.id == *expected_name - && dunder_all_binding(bindings, &entry.id).is_some_and(is_valid_key)) + && bindings + .dunder_all_binding(&entry.id) + .is_some_and(is_valid_key)) .then_some(entry.range) })); } @@ -5731,6 +5722,7 @@ impl<'a> RdepTransaction for Transaction<'a> { definition: &TextRangeWithModule, ) -> Option { self.find_corresponding_interface_definition(sys_info, definition) + .map(|(definition, _)| definition) } fn local_references_from_definition( @@ -5775,6 +5767,7 @@ impl<'a> RdepTransaction for CancellableTransaction<'a> { ) -> Option { self.as_ref() .find_corresponding_interface_definition(sys_info, definition) + .map(|(definition, _)| definition) } fn local_references_from_definition( diff --git a/pyrefly/lib/stubgen/extract.rs b/pyrefly/lib/stubgen/extract.rs index f91e501ddd..fd673b3a11 100644 --- a/pyrefly/lib/stubgen/extract.rs +++ b/pyrefly/lib/stubgen/extract.rs @@ -1021,10 +1021,10 @@ fn resolve_dunder_all(body: &[Stmt], module_info: &Module) -> Option { + DunderAllEntry::Name(_, name, _) => { names.insert(name.clone()); } - DunderAllEntry::Remove(_, name) => { + DunderAllEntry::Remove(_, name, _) => { names.remove(name); } DunderAllEntry::Module(..) => { diff --git a/pyrefly/lib/test/lsp/local_find_refs.rs b/pyrefly/lib/test/lsp/local_find_refs.rs index e3f42dccc0..b9b0b5a1fc 100644 --- a/pyrefly/lib/test/lsp/local_find_refs.rs +++ b/pyrefly/lib/test/lsp/local_find_refs.rs @@ -5,19 +5,13 @@ * LICENSE file in the root directory of this source tree. */ -use std::fs; - use itertools::Itertools; use pretty_assertions::assert_eq; use pyrefly_build::handle::Handle; -use pyrefly_python::module::TextRangeWithModule; 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; @@ -239,63 +233,6 @@ text = "foo" ); } -#[test] -fn dunder_all_global_references_respect_definition_identity() { - let root = TempDir::new().unwrap(); - let mut env = TestEnv::new(); - for (name, code) in [ - ("lib", "foo = 1\n"), - ("exports", "from lib import foo\n__all__ = [\"foo\"]\n"), - ( - "renamed", - "from lib import foo as alias\n__all__ = [\"alias\"]\n", - ), - ("unrelated", "foo = 2\n__all__ = [\"foo\"]\n"), - ] { - let path = root.path().join(format!("{name}.py")); - fs::write(&path, code).unwrap(); - env.add_real_path(name, path); - } - let (state, handle) = env.to_state(); - let handle = handle("lib"); - let mut transaction = state.transaction(); - let definitions = transaction - .find_definition(&handle, TextSize::new(0), FindPreference::default()) - .unwrap(); - let definition = definitions.first(); - let mut references = transaction - .find_global_references_from_definition( - *handle.sys_info(), - definition.metadata.clone(), - TextRangeWithModule::new(definition.module.clone(), definition.definition_range), - ReferenceOptions::all(true), - ) - .unwrap() - .into_iter() - .map(|(module, ranges)| { - ( - module.name().to_string(), - ranges - .into_iter() - .map(|range| module.code_at(range).to_owned()) - .collect::>(), - ) - }) - .collect::>(); - references.sort(); - assert_eq!( - references, - vec![ - ( - "exports".to_owned(), - vec!["foo".to_owned(), "foo".to_owned()] - ), - ("lib".to_owned(), vec!["foo".to_owned()]), - ("renamed".to_owned(), vec!["foo".to_owned()]), - ], - ); -} - // BUG: A __slots__ entry is disconnected from references to its attribute. // Both cursors should find the slots entry, assignment, and read of self.foo (issue #4344). #[test] diff --git a/pyrefly/lib/test/lsp/lsp_interaction/references.rs b/pyrefly/lib/test/lsp/lsp_interaction/references.rs index 4a806f68d1..413c042030 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/references.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/references.rs @@ -45,6 +45,11 @@ fn references_dunder_all_uses_open_buffer_binding() { ])) .unwrap(); } + interaction + .client + .definition("main.py", 2, 13) + .expect_definition_response_from_root("main.py", 1, 0, 1, 3) + .unwrap(); interaction.shutdown().unwrap(); } From 80e232bd001da091b17c647482b6b8cef0d412d8 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 9 Nov 2026 01:33:58 +0900 Subject: [PATCH 09/19] Test renaming from __all__ string entries --- pyrefly/lib/test/lsp/rename.rs | 92 +++++++++++++++++++++++++++++++--- 1 file changed, 85 insertions(+), 7 deletions(-) diff --git a/pyrefly/lib/test/lsp/rename.rs b/pyrefly/lib/test/lsp/rename.rs index 3a055288b3..549f48b1c9 100644 --- a/pyrefly/lib/test/lsp/rename.rs +++ b/pyrefly/lib/test/lsp/rename.rs @@ -361,8 +361,7 @@ fn get_test_report(state: &State, handle: &Handle, position: TextSize) -> String ) } -// BUG: prepare_rename returns None for __all__ and __slots__ entries. -// It should return the range of foo or bar inside the quotes (issue #4344). +// BUG: prepare_rename returns None for the __slots__ entry (issue #4344). #[test] fn prepare_rename_symbol_literals() { let code = r#" @@ -378,17 +377,26 @@ class C: self.bar = 1 "#; let report = get_batched_lsp_operations_report(&[("main", code)], |state, handle, position| { - format!( - "Rename range: {:?}", - state.transaction().prepare_rename(handle, position) - ) + let transaction = state.transaction(); + match transaction.prepare_rename(handle, position) { + Some(range) => { + let module = transaction.get_module_info(handle).unwrap(); + format!( + "Rename range:\n{}", + code_frame_of_source_at_range(module.contents(), range) + ) + } + None => "Rename range: None".to_owned(), + } }); assert_eq!( r#" # main.py 3 | __all__ = ["foo"] ^ -Rename range: None +Rename range: +3 | __all__ = ["foo"] + ^^^ 7 | __slots__ = ("bar",) ^ @@ -399,6 +407,76 @@ Rename range: None ); } +#[test] +fn test_prepare_rename_dunder_all_entries() { + for name in ["foo", "café", "变量"] { + let source = format!( + "{name} = 1\n__all__ = [\"{name}\", r'{name}', \"\"\"{name}\"\"\"]\n__all__.append(\"{name}\")\n__all__.extend([\"{name}\"])\n__all__.remove(\"{name}\")\n" + ); + let (state, handle) = TestEnv::one("main", &source).to_state(); + let handle = handle("main"); + let transaction = state.transaction(); + for (start, text) in source.match_indices(name).skip(1) { + let range = TextRange::at(TextSize::try_from(start).unwrap(), TextSize::of(text)); + for position in [range.start(), range.end()] { + assert_eq!(transaction.prepare_rename(&handle, position), Some(range)); + } + assert_eq!( + rename_local_symbol(&source, start), + source.replace(name, "renamed"), + ); + } + } +} + +#[test] +fn test_prepare_rename_dunder_all_rejects_invalid_targets() { + for (source, literal) in [ + ("foo = 1\n__all__ = [\"foo\"]\ntext = \"foo\"\n", "\"foo\""), + ("__all__ = [\"missing\"]\n", "\"missing\""), + ("foo = 1\n__all__ = [\"f\" \"oo\"]\n", "\"f\""), + ("from math import sin\n__all__ = [\"sin\"]\n", "\"sin\""), + ] { + let (state, handle) = TestEnv::one("main", source).to_state(); + let handle = handle("main"); + let position = TextSize::try_from(source.rfind(literal).unwrap() + 1).unwrap(); + assert_eq!(state.transaction().prepare_rename(&handle, position), None); + } +} + +#[test] +fn test_prepare_rename_dunder_all_import_alias() { + let source = "from math import sin as alias\n__all__ = [\"alias\"]\nalias(0.0)\n"; + let (state, handle) = TestEnv::one("main", source).to_state(); + let handle = handle("main"); + let position = source.find("\"alias\"").unwrap() + 1; + let range = TextRange::at(TextSize::try_from(position).unwrap(), TextSize::of("alias")); + assert_eq!( + state.transaction().prepare_rename(&handle, range.start()), + Some(range), + ); + assert_eq!( + rename_local_symbol(source, position), + source.replace("alias", "renamed"), + ); +} + +#[test] +fn test_rename_dunder_all_updates_source_stub_and_consumers() { + let source = "foo = 1\n__all__ = ['foo']\n"; + let stub = "foo: int\n__all__ = ['foo']\n"; + let consumer = "from lib import foo\n__all__ = ['foo']\nfoo\n"; + let position = source.find("'foo'").unwrap() + 1; + assert_eq!( + rename_paired_symbol(source, stub, consumer, None, position), + vec![ + ("consumer.py".to_owned(), consumer.replace("foo", "renamed")), + ("lib.py".to_owned(), source.replace("foo", "renamed")), + ("lib.pyi".to_owned(), stub.replace("foo", "renamed")), + ], + ); +} + #[test] fn test_rename_parameter_updates_keyword_arguments() { let code = r#" From 374593ffec76af2a240c649a875dbdd180b52568 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 9 Nov 2026 01:38:48 +0900 Subject: [PATCH 10/19] Allow renaming symbols from __all__ entries Use string content ranges to preserve quotes and prefixes. Keep local source definitions during reference operations so rename updates both source files and their stubs. Validate the target before offering rename. --- pyrefly/lib/state/lsp.rs | 38 ++++++++++++++++++++++++++++++++++ pyrefly/lib/test/lsp/rename.rs | 2 ++ 2 files changed, 40 insertions(+) diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 7439c6b0eb..e6f16735ec 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -382,6 +382,18 @@ fn definition_matcher( } } +/// Return the content range of a single string literal with the given expression range. +fn string_literal_content_range(ast: &ModModule, range: TextRange) -> Option { + Ast::locate_node(ast, range.start()) + .into_iter() + .find_map(|node| match node { + AnyNodeRef::ExprStringLiteral(literal) if literal.range() == range => { + Some(literal.as_single_part_string()?.content_range()) + } + _ => None, + }) +} + /// Parses a forward-reference string such as `"list[Foo]"` into the expression the binder /// built from it. Only strings whose source text is exactly their value are parsed: escapes /// or implicit concatenation shift the ranges of the names parsed from a string. @@ -4642,12 +4654,18 @@ impl<'a> Transaction<'a> { pub fn prepare_rename(&self, handle: &Handle, position: TextSize) -> Option { let identifier_context = self.identifier_at(handle, position); + let dunder_all_entry = self.get_exports_data(handle).dunder_all_name_at(position); let definitions = self .find_definition( handle, position, FindPreference { + import_behavior: if dunder_all_entry.is_some() { + ImportBehavior::StopAtRenamedImports + } else { + ImportBehavior::JumpThroughEverything + }, resolve_call_dunders: false, ..Default::default() }, @@ -4655,6 +4673,21 @@ impl<'a> Transaction<'a> { .map(Vec1::into_vec) .unwrap_or_default(); + if dunder_all_entry.is_some() + && (definitions.is_empty() + || definitions.iter().any(|definition| { + let handle = handle + .for_module(definition.module.name(), definition.module.path().dupe()); + definition.module.path().is_bundled() + || self + .get_exports_data(&handle) + .dunder_all_name_at(definition.definition_range.start()) + .is_some() + })) + { + return None; + } + for FindDefinitionItemWithDocstring { module, definition_range, @@ -4687,6 +4720,11 @@ impl<'a> Transaction<'a> { self.setattr_name_literal_at(handle, position) .map(|(name, _)| name.range) }) + .or_else(|| { + let (range, _) = dunder_all_entry?; + let ast = self.get_ast(handle)?; + string_literal_content_range(&ast, range) + }) } pub fn find_local_references( diff --git a/pyrefly/lib/test/lsp/rename.rs b/pyrefly/lib/test/lsp/rename.rs index 549f48b1c9..1403278c3c 100644 --- a/pyrefly/lib/test/lsp/rename.rs +++ b/pyrefly/lib/test/lsp/rename.rs @@ -19,6 +19,7 @@ use ruff_text_size::TextSize; use tempfile::TempDir; use crate::state::lsp::FindPreference; +use crate::state::lsp::ImportBehavior; use crate::state::lsp::ReferenceOptions; use crate::state::require::Require; use crate::state::state::State; @@ -91,6 +92,7 @@ fn replace_paired_references( &source_handle, TextSize::try_from(position).unwrap(), FindPreference { + import_behavior: ImportBehavior::StopAtRenamedImports, resolve_call_dunders: false, ..Default::default() }, From 0fd832e752824ddf091c14264bd7af382fe3adf3 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sat, 14 Nov 2026 13:10:44 +0900 Subject: [PATCH 11/19] Use retained export occurrences when preparing rename --- pyrefly/lib/state/lsp.rs | 54 ++++++++++++++++------------------------ 1 file changed, 22 insertions(+), 32 deletions(-) diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index e6f16735ec..d9f6f40ad5 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -113,6 +113,7 @@ use crate::state::state::Transaction; use crate::state::state::TransactionHandle; use crate::types::module::ModuleType; use crate::types::type_var::Restriction; +use crate::types::types::AnyStyle; use crate::types::types::Type; mod dict_completions; @@ -382,18 +383,6 @@ fn definition_matcher( } } -/// Return the content range of a single string literal with the given expression range. -fn string_literal_content_range(ast: &ModModule, range: TextRange) -> Option { - Ast::locate_node(ast, range.start()) - .into_iter() - .find_map(|node| match node { - AnyNodeRef::ExprStringLiteral(literal) if literal.range() == range => { - Some(literal.as_single_part_string()?.content_range()) - } - _ => None, - }) -} - /// Parses a forward-reference string such as `"list[Foo]"` into the expression the binder /// built from it. Only strings whose source text is exactly their value are parsed: escapes /// or implicit concatenation shift the ranges of the names parsed from a string. @@ -4654,18 +4643,29 @@ impl<'a> Transaction<'a> { pub fn prepare_rename(&self, handle: &Handle, position: TextSize) -> Option { let identifier_context = self.identifier_at(handle, position); - let dunder_all_entry = self.get_exports_data(handle).dunder_all_name_at(position); + let answers = self.get_answers(handle)?; + let bindings = answers.bindings(); + let dunder_all_entry = bindings + .dunder_all_entries() + .iter() + .find(|entry| entry.range.contains_inclusive(position)); + // A string can be renamed only when it identifies an editable, resolved binding. + if let Some(entry) = dunder_all_entry { + let key = bindings.dunder_all_binding(&entry.id)?; + if matches!( + bindings.get(bindings.key_to_idx(key)), + Binding::Any(AnyStyle::Error) + ) { + return None; + } + } let definitions = self .find_definition( handle, position, FindPreference { - import_behavior: if dunder_all_entry.is_some() { - ImportBehavior::StopAtRenamedImports - } else { - ImportBehavior::JumpThroughEverything - }, + import_behavior: ImportBehavior::StopAtRenamedImports, resolve_call_dunders: false, ..Default::default() }, @@ -4675,15 +4675,9 @@ impl<'a> Transaction<'a> { if dunder_all_entry.is_some() && (definitions.is_empty() - || definitions.iter().any(|definition| { - let handle = handle - .for_module(definition.module.name(), definition.module.path().dupe()); - definition.module.path().is_bundled() - || self - .get_exports_data(&handle) - .dunder_all_name_at(definition.definition_range.start()) - .is_some() - })) + || definitions + .iter() + .any(|definition| definition.module.path().is_bundled())) { return None; } @@ -4720,11 +4714,7 @@ impl<'a> Transaction<'a> { self.setattr_name_literal_at(handle, position) .map(|(name, _)| name.range) }) - .or_else(|| { - let (range, _) = dunder_all_entry?; - let ast = self.get_ast(handle)?; - string_literal_content_range(&ast, range) - }) + .or_else(|| dunder_all_entry.map(|entry| entry.range)) } pub fn find_local_references( From ba4af251744d852cfd626b864313840e2ea514a1 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 9 Nov 2026 01:44:26 +0900 Subject: [PATCH 12/19] Test definition lookup from __slots__ entries --- pyrefly/lib/test/lsp/definition.rs | 102 +++++++++++++++++++++++- pyrefly/lib/test/lsp/hover.rs | 2 +- pyrefly/lib/test/lsp/local_find_refs.rs | 7 +- 3 files changed, 106 insertions(+), 5 deletions(-) diff --git a/pyrefly/lib/test/lsp/definition.rs b/pyrefly/lib/test/lsp/definition.rs index d13d36ef81..1db4cf5043 100644 --- a/pyrefly/lib/test/lsp/definition.rs +++ b/pyrefly/lib/test/lsp/definition.rs @@ -3178,7 +3178,6 @@ Definition Result: ); } -// BUG: A __slots__ entry has no definition target. It should resolve to self.foo (issue #4344). #[test] fn dunder_slots_entry_definition_test() { let code = r#" @@ -3195,13 +3194,112 @@ class C: # main.py 3 | __slots__ = ("foo",) ^ -Definition Result: None +Definition Result: +7 | self.foo = 1 + ^^^ "# .trim(), report.trim(), ); } +#[test] +fn dunder_slots_definition_assignment_forms() { + for assignment in [ + "__slots__ = \"foo\"", + "__slots__ = (\"foo\",)", + "__slots__ = [\"foo\"]", + "__slots__ = {\"foo\": \"Field documentation.\"}", + "__slots__: tuple[str, ...] = (\"foo\",)", + "names = __slots__ = (\"foo\",)", + ] { + let code = format!( + "class C:\n {assignment}\n def __init__(self) -> None:\n self.foo: int = 1\n" + ); + let (state, handle) = TestEnv::one("main", &code).to_state(); + let handle = handle("main"); + let position = TextSize::try_from(code.find("\"foo\"").unwrap() + 1).unwrap(); + let definitions = state + .transaction() + .goto_definition(&handle, position) + .unwrap(); + assert_eq!(definitions.len(), 1, "assignment: {assignment}"); + assert_eq!(definitions[0].module.path(), handle.path()); + assert_eq!( + definitions[0].range, + TextRange::at( + TextSize::try_from(code.rfind("foo").unwrap()).unwrap(), + TextSize::of("foo"), + ), + ); + } +} + +#[test] +fn dunder_slots_definition_uses_nearest_class() { + let code = r#" +class Outer: + __slots__ = ("foo",) + def __init__(self) -> None: + self.foo = 1 + + class Inner: + __slots__ = ("foo",) + def __init__(self) -> None: + self.foo = 2 + +def factory(): + class Local: + __slots__ = ("foo",) + def __init__(self) -> None: + self.foo = 3 + return Local +"#; + let (state, handle) = TestEnv::one("main", code).to_state(); + let handle = handle("main"); + let transaction = state.transaction(); + for ((start, _), (definition, _)) in code + .match_indices("\"foo\"") + .zip(code.match_indices("self.foo")) + { + let definitions = transaction + .goto_definition(&handle, TextSize::try_from(start + 1).unwrap()) + .unwrap(); + assert_eq!(definitions.len(), 1); + assert_eq!(definitions[0].module.path(), handle.path()); + assert_eq!( + definitions[0].range, + TextRange::at( + TextSize::try_from(definition + "self.".len()).unwrap(), + TextSize::of("foo"), + ), + ); + } +} + +#[test] +fn dunder_slots_definition_rejects_unrelated_strings() { + for code in [ + "__slots__ = (\"foo\",)\nfoo = 1\n", + "class C:\n def __init__(self):\n __slots__ = (\"foo\",)\n self.foo = 1\n", + "class C:\n __slots__ = {\"foo\": \"foo\"}\n def __init__(self):\n self.foo = 1\n", + "class C:\n names = (\"foo\",)\n def __init__(self):\n self.foo = 1\n", + "class C:\n __slots__ = tuple([\"foo\"])\n def __init__(self):\n self.foo = 1\n", + "class C:\n __slots__ = (\"foo\",)\n", + ] { + let (state, handle) = TestEnv::one("main", code).to_state(); + let handle = handle("main"); + let position = TextSize::try_from(code.rfind("\"foo\"").unwrap() + 1).unwrap(); + assert!( + state + .transaction() + .goto_definition(&handle, position) + .is_err(), + "code: {code}", + ); + } +} + #[test] fn string_literal_not_in_dunder_all() { let pkg = r#" diff --git a/pyrefly/lib/test/lsp/hover.rs b/pyrefly/lib/test/lsp/hover.rs index b2f9651d8f..ba86037f09 100644 --- a/pyrefly/lib/test/lsp/hover.rs +++ b/pyrefly/lib/test/lsp/hover.rs @@ -78,7 +78,7 @@ class C: 7 | __slots__ = ("bar",) ^ ```python -Literal['bar'] +(attribute) bar: Literal['bar'] ``` "# .trim(), diff --git a/pyrefly/lib/test/lsp/local_find_refs.rs b/pyrefly/lib/test/lsp/local_find_refs.rs index b9b0b5a1fc..fa70064922 100644 --- a/pyrefly/lib/test/lsp/local_find_refs.rs +++ b/pyrefly/lib/test/lsp/local_find_refs.rs @@ -233,7 +233,7 @@ text = "foo" ); } -// BUG: A __slots__ entry is disconnected from references to its attribute. +// BUG: References omit __slots__ entries. // Both cursors should find the slots entry, assignment, and read of self.foo (issue #4344). #[test] fn dunder_slots_entries_are_attribute_references() { @@ -258,7 +258,10 @@ class C: 3 | __slots__ = ("foo",) ^ References: - +7 | self.foo = 1 + ^^^ +11 | return self.foo + ^^^ 7 | self.foo = 1 ^ From 6e346ddb218b6c82d203397280601de459438358 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 9 Nov 2026 01:46:40 +0900 Subject: [PATCH 13/19] Resolve __slots__ strings to their attribute definitions Read the owning class and literal slot name from the existing AST nodes. Use attribute definition lookup to navigate to the field declaration. --- pyrefly/lib/state/lsp.rs | 61 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 61 insertions(+) diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index d9f6f40ad5..2952e86a83 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -383,6 +383,56 @@ fn definition_matcher( } } +/// Identify a literal slot name and the class definition that owns it. +fn dunder_slots_entry(covering_nodes: &[AnyNodeRef]) -> Option<(ShortIdentifier, Name)> { + let literal = covering_nodes.iter().find_map(|node| match node { + AnyNodeRef::ExprStringLiteral(literal) => Some(literal), + _ => None, + })?; + let value = covering_nodes.iter().find_map(|node| match node { + AnyNodeRef::StmtAssign(assign) + if assign + .targets + .iter() + .any(|target| matches!(target, Expr::Name(name) if name.id == *dunder::SLOTS)) => + { + Some(assign.value.as_ref()) + } + AnyNodeRef::StmtAnnAssign(assign) + if matches!(assign.target.as_ref(), Expr::Name(name) if name.id == *dunder::SLOTS) => + { + assign.value.as_deref() + } + _ => None, + })?; + let is_entry = match value { + Expr::StringLiteral(value) => value.range() == literal.range(), + Expr::Tuple(tuple) => tuple.elts.iter().any(|elt| elt.range() == literal.range()), + Expr::List(list) => list.elts.iter().any(|elt| elt.range() == literal.range()), + Expr::Dict(dict) => dict.items.iter().any(|item| { + item.key + .as_ref() + .is_some_and(|key| key.range() == literal.range()) + }), + _ => false, + }; + if !is_entry { + return None; + } + let Some(AnyNodeRef::StmtClassDef(class)) = covering_nodes.iter().find(|node| { + matches!( + node, + AnyNodeRef::StmtClassDef(_) | AnyNodeRef::StmtFunctionDef(_) + ) + }) else { + return None; + }; + Some(( + ShortIdentifier::new(&class.name), + Name::new(literal.value.to_str()), + )) +} + /// Parses a forward-reference string such as `"list[Foo]"` into the expression the binder /// built from it. Only strings whose source text is exactly their value are parsed: escapes /// or implicit concatenation shift the ranges of the names parsed from a string. @@ -3120,6 +3170,17 @@ impl<'a> Transaction<'a> { { return Ok(vec1![definition]); } + if let Some((class, name)) = dunder_slots_entry(&covering_nodes) + && let Some(answers) = self.get_answers(handle) + && let Some(idx) = answers + .bindings() + .key_to_idx_hashed_opt(Hashed::new(&Key::Definition(class))) + && let Some(class_type) = answers.get_type_at(idx) + { + return self.find_attribute_definition_for_base_type( + handle, preference, class_type, &name, + ); + } } if let Some(AnyNodeRef::ExprStringLiteral(literal)) = covering_nodes From c7d6e7a26d1c2831961e304ffef2608f9aefd0a9 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 9 Nov 2026 01:50:28 +0900 Subject: [PATCH 14/19] Test hover types for __slots__ entries --- pyrefly/lib/test/lsp/hover.rs | 78 ++++++++++++++++++++++++++++++++++- 1 file changed, 76 insertions(+), 2 deletions(-) diff --git a/pyrefly/lib/test/lsp/hover.rs b/pyrefly/lib/test/lsp/hover.rs index ba86037f09..b3ff93788f 100644 --- a/pyrefly/lib/test/lsp/hover.rs +++ b/pyrefly/lib/test/lsp/hover.rs @@ -42,7 +42,6 @@ fn get_test_report(state: &State, handle: &Handle, position: TextSize) -> String } } -// BUG: Hover shows a string literal type for the __slots__ entry (issue #4344). #[test] fn hover_symbol_literals() { let code = r#" @@ -78,7 +77,7 @@ class C: 7 | __slots__ = ("bar",) ^ ```python -(attribute) bar: Literal['bar'] +(attribute) bar: str ``` "# .trim(), @@ -86,6 +85,81 @@ class C: ); } +#[test] +fn hover_dunder_slots_assignment_forms() { + for assignment in [ + "__slots__ = \"foo\"", + "__slots__ = (\"foo\",)", + "__slots__ = [\"foo\"]", + "__slots__ = {\"foo\": \"Field documentation.\"}", + "__slots__: tuple[str, ...] = (\"foo\",)", + "names = __slots__ = (\"foo\",)", + ] { + let code = format!( + "class C:\n {assignment}\n def __init__(self) -> None:\n self.foo: int = 1\n" + ); + let (state, handle) = TestEnv::one("main", &code).to_state(); + let handle = handle("main"); + let position = TextSize::try_from(code.find("\"foo\"").unwrap() + 1).unwrap(); + let report = get_test_report(&state, &handle, position); + assert!(report.contains("(attribute) foo: int"), "got: {report}"); + } +} + +#[test] +fn hover_dunder_slots_preserves_type_parameters() { + let code = r#" +class C[T]: + __slots__ = ("foo",) +# ^ + def __init__(self, value: T) -> None: + self.foo: T = value +"#; + let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report); + assert!(report.contains("(attribute) foo: T"), "got: {report}"); +} + +#[test] +fn hover_dunder_slots_uses_nearest_class() { + let code = r#" +class Outer: + __slots__ = ("foo",) + def __init__(self) -> None: + self.foo: int = 1 + + class Inner: + __slots__ = ("foo",) + def __init__(self) -> None: + self.foo: str = "" +"#; + let (state, handle) = TestEnv::one("main", code).to_state(); + let handle = handle("main"); + for ((start, _), expected_type) in code.match_indices("\"foo\"").zip(["int", "str"]) { + let position = TextSize::try_from(start + 1).unwrap(); + let report = get_test_report(&state, &handle, position); + assert!( + report.contains(&format!("(attribute) foo: {expected_type}")), + "got: {report}", + ); + } +} + +#[test] +fn hover_dunder_slots_keeps_unresolved_and_unrelated_literal_types() { + for code in [ + "class C:\n __slots__ = (\"foo\",)\n", + "class C:\n __slots__ = {\"foo\": \"foo\"}\n def __init__(self):\n self.foo: int = 1\n", + "class C:\n def __init__(self):\n __slots__ = (\"foo\",)\n self.foo: int = 1\n", + "class C:\n names = (\"foo\",)\n def __init__(self):\n self.foo: int = 1\n", + ] { + let (state, handle) = TestEnv::one("main", code).to_state(); + let handle = handle("main"); + let position = TextSize::try_from(code.rfind("\"foo\"").unwrap() + 1).unwrap(); + let report = get_test_report(&state, &handle, position); + assert!(report.contains("Literal['foo']"), "got: {report}"); + } +} + #[test] fn hover_dunder_all_reexport_entries() { let code = r#" From f3f2bf2ac51f7c5f05db3d4689ca67d995ccfbea Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 9 Nov 2026 01:53:05 +0900 Subject: [PATCH 15/19] Show attribute types when hovering over __slots__ entries Resolve slot names against a generic instance of their owning class. This displays the field type while preserving its type parameters. --- pyrefly/lib/lsp/wasm/hover.rs | 5 +++-- pyrefly/lib/state/lsp.rs | 24 ++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/pyrefly/lib/lsp/wasm/hover.rs b/pyrefly/lib/lsp/wasm/hover.rs index 2fa7c2c0a5..383a8b7123 100644 --- a/pyrefly/lib/lsp/wasm/hover.rs +++ b/pyrefly/lib/lsp/wasm/hover.rs @@ -1402,8 +1402,9 @@ fn resolve_hovered_type( ast: Option<&ModModule>, position: TextSize, ) -> Option { - let mut type_ = transaction - .subscript_operator_type_at(handle, position) + let mut type_ = ast + .and_then(|ast| transaction.dunder_slots_entry_type(handle, ast, position)) + .or_else(|| transaction.subscript_operator_type_at(handle, position)) .or_else(|| transaction.get_type_at(handle, position)) .or_else(|| transaction.operator_type_at(handle, position))?; diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 2952e86a83..dfd8d448d6 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -3014,6 +3014,30 @@ impl<'a> Transaction<'a> { None } + /// Return the instance attribute type named by a literal `__slots__` entry. + pub(crate) fn dunder_slots_entry_type( + &self, + handle: &Handle, + ast: &ModModule, + position: TextSize, + ) -> Option { + let (class, name) = dunder_slots_entry(&Ast::locate_node(ast, position))?; + let answers = self.get_answers(handle)?; + let idx = answers + .bindings() + .key_to_idx_hashed_opt(Hashed::new(&Key::Definition(class)))?; + let Type::ClassDef(class) = answers.get_type_at(idx)? else { + return None; + }; + self.ad_hoc_solve(handle, "dunder_slots_entry_type", |solver| { + solver + .completions(solver.instantiate(&class), Some(&name), true, false) + .into_iter() + .find(|attribute| attribute.name == name) + .and_then(|attribute| attribute.ty) + })? + } + fn find_definition_for_keyword_argument( &self, handle: &Handle, From cbb8633ae8ced07eb3413095847ebf5a44da05eb Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sat, 14 Nov 2026 13:14:41 +0900 Subject: [PATCH 16/19] Share checked slots entries across definition and hover --- pyrefly/lib/alt/class/class_metadata.rs | 30 +------ pyrefly/lib/binding/binding.rs | 32 +++++++ pyrefly/lib/binding/bindings.rs | 45 +++++++++- pyrefly/lib/lsp/wasm/hover.rs | 5 +- pyrefly/lib/state/lsp.rs | 107 +++++++----------------- pyrefly/lib/test/lsp/definition.rs | 2 +- pyrefly/lib/test/lsp/hover.rs | 2 +- 7 files changed, 113 insertions(+), 110 deletions(-) diff --git a/pyrefly/lib/alt/class/class_metadata.rs b/pyrefly/lib/alt/class/class_metadata.rs index 2d106a37f3..6503fec2c4 100644 --- a/pyrefly/lib/alt/class/class_metadata.rs +++ b/pyrefly/lib/alt/class/class_metadata.rs @@ -5,8 +5,6 @@ * LICENSE file in the root directory of this source tree. */ -use std::iter; - use dupe::Dupe; use dupe::IterDupedExt; use itertools::Itertools; @@ -69,7 +67,6 @@ use crate::binding::base_class::BaseClassGeneric; use crate::binding::base_class::BaseClassGenericKind; use crate::binding::binding::Binding; use crate::binding::binding::ClassFieldDefinition; -use crate::binding::binding::ExprOrBinding; use crate::binding::binding::Key; use crate::binding::binding::KeyAnnotation; use crate::binding::binding::KeyClassField; @@ -710,33 +707,10 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { return ExplicitSlots::Absent; }; let binding = self.bindings().get::(idx); - let ClassFieldDefinition::AssignedInBody { value, .. } = &binding.definition else { - return ExplicitSlots::Unknown; - }; - let ExprOrBinding::Expr(expr) = value.as_ref() else { - return ExplicitSlots::Unknown; - }; - - fn name_of(expr: &Expr) -> Option { - match expr { - Expr::StringLiteral(s) => Some(Name::new(s.value.to_str())), - _ => None, - } - } - let names: Option> = match expr { - Expr::Tuple(t) => t.elts.iter().map(name_of).collect(), - Expr::List(l) => l.elts.iter().map(name_of).collect(), - Expr::Dict(d) => d - .items - .iter() - .map(|item| item.key.as_ref().and_then(name_of)) - .collect(), - Expr::StringLiteral(s) => Some(iter::once(Name::new(s.value.to_str())).collect()), - _ => return ExplicitSlots::Unknown, - }; - let Some(names) = names else { + let Some(entries) = binding.explicit_slots() else { return ExplicitSlots::Unknown; }; + let names: SmallSet = entries.into_iter().map(|(name, _)| name).collect(); // A metaclass other than `type` or `ABCMeta` may remove names from the namespace before it // creates the class. if cls.module().path().is_interface() diff --git a/pyrefly/lib/binding/binding.rs b/pyrefly/lib/binding/binding.rs index 2520235c43..fa020b06c9 100644 --- a/pyrefly/lib/binding/binding.rs +++ b/pyrefly/lib/binding/binding.rs @@ -3727,6 +3727,38 @@ pub struct BindingClassField { pub definition: ClassFieldDefinition, } +impl BindingClassField { + /// Extract literal slot names and their editable content ranges from the checked field. + pub fn explicit_slots(&self) -> Option)>> { + let ClassFieldDefinition::AssignedInBody { value, .. } = &self.definition else { + return None; + }; + let ExprOrBinding::Expr(expr) = value.as_ref() else { + return None; + }; + let name_of = |expr: &Expr| match expr { + Expr::StringLiteral(literal) => Some(( + Name::new(literal.value.to_str()), + literal + .as_single_part_string() + .map(|part| part.content_range()), + )), + _ => None, + }; + match expr { + Expr::Tuple(tuple) => tuple.elts.iter().map(name_of).collect(), + Expr::List(list) => list.elts.iter().map(name_of).collect(), + Expr::Dict(dict) => dict + .items + .iter() + .map(|item| item.key.as_ref().and_then(name_of)) + .collect(), + Expr::StringLiteral(_) => Some(vec![name_of(expr)?]), + _ => None, + } + } +} + impl DisplayWith for BindingClassField { fn fmt(&self, f: &mut fmt::Formatter<'_>, ctx: &Bindings) -> fmt::Result { write!( diff --git a/pyrefly/lib/binding/bindings.rs b/pyrefly/lib/binding/bindings.rs index e11472d16f..e1d72a606f 100644 --- a/pyrefly/lib/binding/bindings.rs +++ b/pyrefly/lib/binding/bindings.rs @@ -417,6 +417,13 @@ table! { pub struct BindingTable(pub BindingEntry) } +/// A literal slot occurrence and the class whose field binding declares it. +#[derive(Debug)] +pub struct SlotEntry { + pub name: Identifier, + pub class_idx: Idx, +} + #[derive(Debug)] pub struct Bindings { module_info: ModuleInfo, @@ -435,6 +442,7 @@ pub struct Bindings { docstring_type_ranges: Vec, /// Each entry records an `__all__` name and its range without string delimiters. dunder_all_entries: Vec, + dunder_slots_entries: Vec, scope_trace: Option, /// Names `del`eted at module scope; retained outside the LSP `scope_trace`. module_deletes: SmallSet, @@ -697,6 +705,7 @@ impl Bindings { docstring_type_ranges: Vec::new(), scope_trace: None, dunder_all_entries: Vec::new(), + dunder_slots_entries: Vec::new(), module_deletes: SmallSet::new(), unset_nodes: UnsetNodes::default(), global_declarations: SmallSet::new(), @@ -954,6 +963,16 @@ impl Bindings { } } + pub fn dunder_slots_entries(&self) -> &[SlotEntry] { + &self.dunder_slots_entries + } + + pub fn dunder_slots_entry_at(&self, position: TextSize) -> Option<&SlotEntry> { + self.dunder_slots_entries + .iter() + .find(|entry| entry.name.range.contains_inclusive(position)) + } + pub fn dunder_all_entries(&self) -> &[Identifier] { &self.dunder_all_entries } @@ -1306,7 +1325,7 @@ impl Bindings { classes: builder.django_relation_classes.into_boxed_slice(), }, ); - Self { + let mut bindings = Self { module_info, parse_errors, binding_errors: builder.errors.finish(), @@ -1317,6 +1336,7 @@ impl Bindings { annotated_metadata_ranges: merge_ranges(builder.annotated_metadata_ranges), docstring_type_ranges: merge_ranges(builder.docstring_type_ranges), dunder_all_entries, + dunder_slots_entries: Vec::new(), scope_trace: if enable_trace { Some(scope_trace) } else { @@ -1341,7 +1361,28 @@ impl Bindings { subsequently_initialized: builder.subsequently_initialized, promote_ranges: builder.promote_ranges, functools_wrapper_mutation, - } + }; + bindings.dunder_slots_entries = bindings + .keys::() + .flat_map(|idx| { + let field = bindings.get(idx); + if field.name != *dunder::SLOTS { + return Vec::new(); + } + field + .explicit_slots() + .into_iter() + .flatten() + .filter_map(|(name, range)| { + Some(SlotEntry { + name: Identifier::new(name, range?), + class_idx: field.class_idx, + }) + }) + .collect() + }) + .collect(); + bindings } fn should_emit_semantic_syntax_error(error: &SemanticSyntaxError) -> bool { diff --git a/pyrefly/lib/lsp/wasm/hover.rs b/pyrefly/lib/lsp/wasm/hover.rs index 383a8b7123..2fa7c2c0a5 100644 --- a/pyrefly/lib/lsp/wasm/hover.rs +++ b/pyrefly/lib/lsp/wasm/hover.rs @@ -1402,9 +1402,8 @@ fn resolve_hovered_type( ast: Option<&ModModule>, position: TextSize, ) -> Option { - let mut type_ = ast - .and_then(|ast| transaction.dunder_slots_entry_type(handle, ast, position)) - .or_else(|| transaction.subscript_operator_type_at(handle, position)) + let mut type_ = transaction + .subscript_operator_type_at(handle, position) .or_else(|| transaction.get_type_at(handle, position)) .or_else(|| transaction.operator_type_at(handle, position))?; diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index dfd8d448d6..f63892e634 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -383,56 +383,6 @@ fn definition_matcher( } } -/// Identify a literal slot name and the class definition that owns it. -fn dunder_slots_entry(covering_nodes: &[AnyNodeRef]) -> Option<(ShortIdentifier, Name)> { - let literal = covering_nodes.iter().find_map(|node| match node { - AnyNodeRef::ExprStringLiteral(literal) => Some(literal), - _ => None, - })?; - let value = covering_nodes.iter().find_map(|node| match node { - AnyNodeRef::StmtAssign(assign) - if assign - .targets - .iter() - .any(|target| matches!(target, Expr::Name(name) if name.id == *dunder::SLOTS)) => - { - Some(assign.value.as_ref()) - } - AnyNodeRef::StmtAnnAssign(assign) - if matches!(assign.target.as_ref(), Expr::Name(name) if name.id == *dunder::SLOTS) => - { - assign.value.as_deref() - } - _ => None, - })?; - let is_entry = match value { - Expr::StringLiteral(value) => value.range() == literal.range(), - Expr::Tuple(tuple) => tuple.elts.iter().any(|elt| elt.range() == literal.range()), - Expr::List(list) => list.elts.iter().any(|elt| elt.range() == literal.range()), - Expr::Dict(dict) => dict.items.iter().any(|item| { - item.key - .as_ref() - .is_some_and(|key| key.range() == literal.range()) - }), - _ => false, - }; - if !is_entry { - return None; - } - let Some(AnyNodeRef::StmtClassDef(class)) = covering_nodes.iter().find(|node| { - matches!( - node, - AnyNodeRef::StmtClassDef(_) | AnyNodeRef::StmtFunctionDef(_) - ) - }) else { - return None; - }; - Some(( - ShortIdentifier::new(&class.name), - Name::new(literal.value.to_str()), - )) -} - /// Parses a forward-reference string such as `"list[Foo]"` into the expression the binder /// built from it. Only strings whose source text is exactly their value are parsed: escapes /// or implicit concatenation shift the ranges of the names parsed from a string. @@ -1552,6 +1502,9 @@ impl<'a> Transaction<'a> { .get_hashed_opt(Hashed::new(&KeyExport(name))) .cloned(); } + if let Some(ty) = self.dunder_slots_entry_type(handle, position) { + return Some(ty); + } return self .type_from_match_wildcard_at_impl(handle, position) .or_else(|| self.type_from_expression_at_impl(handle, position, false)); @@ -3015,26 +2968,24 @@ impl<'a> Transaction<'a> { } /// Return the instance attribute type named by a literal `__slots__` entry. - pub(crate) fn dunder_slots_entry_type( - &self, - handle: &Handle, - ast: &ModModule, - position: TextSize, - ) -> Option { - let (class, name) = dunder_slots_entry(&Ast::locate_node(ast, position))?; + fn dunder_slots_entry_type(&self, handle: &Handle, position: TextSize) -> Option { let answers = self.get_answers(handle)?; - let idx = answers - .bindings() - .key_to_idx_hashed_opt(Hashed::new(&Key::Definition(class)))?; - let Type::ClassDef(class) = answers.get_type_at(idx)? else { - return None; - }; + let entry = answers.bindings().dunder_slots_entry_at(position)?; self.ad_hoc_solve(handle, "dunder_slots_entry_type", |solver| { - solver - .completions(solver.instantiate(&class), Some(&name), true, false) - .into_iter() - .find(|attribute| attribute.name == name) - .and_then(|attribute| attribute.ty) + let class = solver.get_idx(entry.class_idx).0.as_ref()?; + let instance = solver.instantiate(class); + if !solver.has_attr(&instance, &entry.name.id) { + return None; + } + Some(solver.type_of_attr_get( + &instance, + &entry.name.id, + entry.name.range, + &solver.error_swallower(), + ErrorKind::MissingAttribute, + None, + "dunder_slots_entry_type", + )) })? } @@ -3194,15 +3145,21 @@ impl<'a> Transaction<'a> { { return Ok(vec1![definition]); } - if let Some((class, name)) = dunder_slots_entry(&covering_nodes) - && let Some(answers) = self.get_answers(handle) - && let Some(idx) = answers - .bindings() - .key_to_idx_hashed_opt(Hashed::new(&Key::Definition(class))) - && let Some(class_type) = answers.get_type_at(idx) + if let Some(answers) = self.get_answers(handle) + && let Some(entry) = answers.bindings().dunder_slots_entry_at(position) + && let Some(class_type) = self + .ad_hoc_solve(handle, "dunder_slots_definition", |solver| { + Some(Type::ClassDef( + solver.get_idx(entry.class_idx).0.as_ref()?.dupe(), + )) + }) + .flatten() { return self.find_attribute_definition_for_base_type( - handle, preference, class_type, &name, + handle, + preference, + class_type, + &entry.name.id, ); } } diff --git a/pyrefly/lib/test/lsp/definition.rs b/pyrefly/lib/test/lsp/definition.rs index 1db4cf5043..f71ac994e1 100644 --- a/pyrefly/lib/test/lsp/definition.rs +++ b/pyrefly/lib/test/lsp/definition.rs @@ -3211,7 +3211,6 @@ fn dunder_slots_definition_assignment_forms() { "__slots__ = [\"foo\"]", "__slots__ = {\"foo\": \"Field documentation.\"}", "__slots__: tuple[str, ...] = (\"foo\",)", - "names = __slots__ = (\"foo\",)", ] { let code = format!( "class C:\n {assignment}\n def __init__(self) -> None:\n self.foo: int = 1\n" @@ -3283,6 +3282,7 @@ fn dunder_slots_definition_rejects_unrelated_strings() { "__slots__ = (\"foo\",)\nfoo = 1\n", "class C:\n def __init__(self):\n __slots__ = (\"foo\",)\n self.foo = 1\n", "class C:\n __slots__ = {\"foo\": \"foo\"}\n def __init__(self):\n self.foo = 1\n", + "class C:\n names = __slots__ = (\"foo\",)\n def __init__(self):\n self.foo: int = 1\n", "class C:\n names = (\"foo\",)\n def __init__(self):\n self.foo = 1\n", "class C:\n __slots__ = tuple([\"foo\"])\n def __init__(self):\n self.foo = 1\n", "class C:\n __slots__ = (\"foo\",)\n", diff --git a/pyrefly/lib/test/lsp/hover.rs b/pyrefly/lib/test/lsp/hover.rs index b3ff93788f..fa62b24764 100644 --- a/pyrefly/lib/test/lsp/hover.rs +++ b/pyrefly/lib/test/lsp/hover.rs @@ -93,7 +93,6 @@ fn hover_dunder_slots_assignment_forms() { "__slots__ = [\"foo\"]", "__slots__ = {\"foo\": \"Field documentation.\"}", "__slots__: tuple[str, ...] = (\"foo\",)", - "names = __slots__ = (\"foo\",)", ] { let code = format!( "class C:\n {assignment}\n def __init__(self) -> None:\n self.foo: int = 1\n" @@ -150,6 +149,7 @@ fn hover_dunder_slots_keeps_unresolved_and_unrelated_literal_types() { "class C:\n __slots__ = (\"foo\",)\n", "class C:\n __slots__ = {\"foo\": \"foo\"}\n def __init__(self):\n self.foo: int = 1\n", "class C:\n def __init__(self):\n __slots__ = (\"foo\",)\n self.foo: int = 1\n", + "class C:\n names = __slots__ = (\"foo\",)\n def __init__(self):\n self.foo: int = 1\n", "class C:\n names = (\"foo\",)\n def __init__(self):\n self.foo: int = 1\n", ] { let (state, handle) = TestEnv::one("main", code).to_state(); From bb888342d846449a80e21495be4350bcbcfbc222 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 9 Nov 2026 02:00:25 +0900 Subject: [PATCH 17/19] Test references in __slots__ string entries --- pyrefly/lib/test/lsp/local_find_refs.rs | 170 +++++++++++++++++++++++- 1 file changed, 168 insertions(+), 2 deletions(-) diff --git a/pyrefly/lib/test/lsp/local_find_refs.rs b/pyrefly/lib/test/lsp/local_find_refs.rs index fa70064922..56b3dbd1ce 100644 --- a/pyrefly/lib/test/lsp/local_find_refs.rs +++ b/pyrefly/lib/test/lsp/local_find_refs.rs @@ -5,13 +5,22 @@ * LICENSE file in the root directory of this source tree. */ +use std::fs; +use std::slice; + use itertools::Itertools; use pretty_assertions::assert_eq; use pyrefly_build::handle::Handle; +use pyrefly_python::module::TextRangeWithModule; +use ruff_text_size::TextRange; use ruff_text_size::TextSize; +use tempfile::TempDir; +use crate::state::lsp::FindPreference; use crate::state::lsp::ReferenceOptions; +use crate::state::require::Require; 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; @@ -233,8 +242,6 @@ text = "foo" ); } -// BUG: References omit __slots__ entries. -// Both cursors should find the slots entry, assignment, and read of self.foo (issue #4344). #[test] fn dunder_slots_entries_are_attribute_references() { let code = r#" @@ -258,6 +265,8 @@ class C: 3 | __slots__ = ("foo",) ^ References: +3 | __slots__ = ("foo",) + ^^^ 7 | self.foo = 1 ^^^ 11 | return self.foo @@ -266,6 +275,8 @@ References: 7 | self.foo = 1 ^ References: +3 | __slots__ = ("foo",) + ^^^ 7 | self.foo = 1 ^^^ 11 | return self.foo @@ -276,6 +287,161 @@ References: ); } +#[test] +fn dunder_slots_references_use_string_contents() { + for name in ["foo", "café", "变量"] { + for assignment in [ + "__slots__ = \"foo\"", + "__slots__ = (\"foo\",)", + "__slots__ = [r'foo']", + "__slots__ = {\"foo\": \"Field documentation.\"}", + "__slots__: tuple[str, ...] = (\"\"\"foo\"\"\",)", + "names = __slots__ = (\"foo\",)", + ] { + let assignment = assignment.replace("foo", name); + let code = format!( + "class C:\n {assignment}\n {name}: int\n def read(self) -> int:\n return self.{name}\n text = \"{name}\"\n" + ); + let (state, handle) = TestEnv::one("main", &code).to_state(); + let handle = handle("main"); + let transaction = state.transaction(); + let module = transaction.get_module_info(&handle).unwrap(); + for position in [ + code.find(name).unwrap(), + code.find(&format!("self.{name}")).unwrap() + "self.".len(), + ] { + for include_declaration in [false, true] { + let references = transaction.find_local_references( + &handle, + TextSize::try_from(position).unwrap(), + ReferenceOptions::all(include_declaration), + ); + let expected_lines = if include_declaration { + vec![1, 2, 4] + } else { + vec![1, 4] + }; + assert_eq!( + references + .iter() + .map(|range| module.to_lsp_range(*range).start.line) + .collect::>(), + expected_lines, + "code: {code}", + ); + for range in references { + assert_eq!(module.code_at(range), name); + } + } + } + } + } +} + +#[test] +fn dunder_slots_references_respect_class_identity() { + let code = r#" +class C: + __slots__ = {"foo": "foo"} + text = "foo" + def __init__(self) -> None: + __slots__ = ("foo",) + self.foo = 1 + def read(self) -> int: + return self.foo + + class Inner: + __slots__ = ("foo",) + foo: str + +class Other: + __slots__ = ("foo",) + foo: str +"#; + let (state, handle) = TestEnv::one("main", code).to_state(); + let handle = handle("main"); + let expected = [ + code.find("\"foo\"").unwrap() + 1, + code.find("self.foo").unwrap() + "self.".len(), + code.rfind("self.foo").unwrap() + "self.".len(), + ] + .into_iter() + .map(|start| TextRange::at(TextSize::try_from(start).unwrap(), TextSize::of("foo"))) + .collect::>(); + for range in &expected { + assert_eq!( + state.transaction().find_local_references( + &handle, + range.start(), + ReferenceOptions::all(true), + ), + expected, + ); + } +} + +#[test] +fn dunder_slots_global_references_in_indexed_files() { + let root = TempDir::new().unwrap(); + let mut env = TestEnv::new() + .with_default_require_level(Require::Indexing) + .with_run_require(Require::Indexing); + let base = "class Base:\n foo: int\n"; + for (name, code) in [ + ("base", base), + ( + "child", + "from base import Base\nclass Child(Base):\n __slots__ = (\"foo\",)\n", + ), + ( + "unrelated", + "from base import Base\nclass Other:\n __slots__ = (\"foo\",)\n foo: str\n", + ), + ("plain", "from base import Base\nclass Plain(Base): pass\n"), + ] { + let path = root.path().join(format!("{name}.py")); + fs::write(&path, code).unwrap(); + env.add_real_path(name, path); + } + let (state, handle) = env.to_state(); + let base_handle = handle("base"); + let child_handle = handle("child"); + let mut transaction = state.transaction(); + assert!(transaction.get_ast(&child_handle).is_none()); + assert!(transaction.get_answers(&child_handle).is_none()); + transaction.run(slice::from_ref(&base_handle), Require::Everything, None); + let definitions = transaction + .find_definition( + &base_handle, + TextSize::try_from(base.find("foo").unwrap()).unwrap(), + FindPreference::default(), + ) + .unwrap(); + let definition = definitions.first(); + let mut references = transaction + .find_global_references_from_definition( + *base_handle.sys_info(), + definition.metadata.clone(), + TextRangeWithModule::new(definition.module.clone(), definition.definition_range), + ReferenceOptions::all(true), + ) + .unwrap() + .into_iter() + .map(|(module, ranges)| { + for range in &ranges { + assert_eq!(module.code_at(*range), "foo"); + } + (module.name().to_string(), ranges.len()) + }) + .collect::>(); + references.sort(); + assert_eq!( + references, + vec![("base".to_owned(), 1), ("child".to_owned(), 1)] + ); + assert!(transaction.get_answers(&handle("plain")).is_none()); +} + #[test] fn references_to_aliased_submodule_stop_at_import() { let main = r#" From 3b9bb52aa1749f3351e117dd9396caa581c93984 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 9 Nov 2026 02:03:56 +0900 Subject: [PATCH 18/19] Include __slots__ strings in attribute reference results Resolve each slot string to its attribute before adding its content range. Load class bindings only for candidate modules with __slots__ fields so indexed files also return these references. --- pyrefly/lib/binding/metadata.rs | 6 +++ pyrefly/lib/state/lsp.rs | 68 ++++++++++++++++++++++++++++++--- 2 files changed, 69 insertions(+), 5 deletions(-) diff --git a/pyrefly/lib/binding/metadata.rs b/pyrefly/lib/binding/metadata.rs index d1254a9a0d..c31396c038 100644 --- a/pyrefly/lib/binding/metadata.rs +++ b/pyrefly/lib/binding/metadata.rs @@ -13,6 +13,7 @@ use pyrefly_types::class::ClassDefIndex; use pyrefly_types::class::ClassFields; +use ruff_python_ast::name::Name; /// Metadata for a single class definition, populated during binding. #[derive(Debug, Clone, Default)] @@ -67,4 +68,9 @@ impl BindingsMetadata { pub fn get_class_mut(&mut self, idx: ClassDefIndex) -> &mut ClassMetadata { &mut self.classes[idx.0 as usize] } + + /// Return whether any class in this module defines a field with this name. + pub fn has_class_field(&self, name: &Name) -> bool { + self.classes.iter().any(|class| class.fields.contains(name)) + } } diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index f63892e634..2341460b2c 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -4950,6 +4950,10 @@ impl<'a> Transaction<'a> { let definition_name = Name::new(module.code_at(definition_range)); let is_parameter_definition = definition_metadata.symbol_kind() == Some(SymbolKind::Parameter); + let is_attribute = matches!( + definition_metadata, + DefinitionMetadata::Attribute | DefinitionMetadata::VariableOrAttribute(_) + ); let mut references = if handle.path() != module.path() { self.local_references_from_external_definition(handle, definition_range, module)? } else { @@ -4961,11 +4965,7 @@ impl<'a> Transaction<'a> { options.include_declaration, )? }; - if matches!( - definition_metadata, - DefinitionMetadata::Attribute | DefinitionMetadata::VariableOrAttribute(_) - ) && let Some(answers) = self.get_answers(handle) - { + if is_attribute && let Some(answers) = self.get_answers(handle) { let literals = answers .bindings() .setattr_name_literals() @@ -4999,6 +4999,47 @@ impl<'a> Transaction<'a> { &definition_name, )); } + if is_attribute + && self + .get_bindings_metadata(handle) + .is_some_and(|metadata| metadata.has_class_field(&dunder::SLOTS)) + && let Some(ast) = self.get_ast(handle) + { + fn collect(expr: &Expr, name: &Name, ranges: &mut Vec) { + if let Expr::StringLiteral(literal) = expr + && literal.value.to_str() == name.as_str() + && let Some(part) = literal.as_single_part_string() + { + ranges.push(part.content_range()); + } + expr.recurse(&mut |expr| collect(expr, name, ranges)); + } + let mut literals = Vec::new(); + ast.visit(&mut |expr| collect(expr, &definition_name, &mut literals)); + let matches_definition = definition_matcher(module, definition_range); + for range in literals { + if dunder_slots_entry(&Ast::locate_node(&ast, range.start())).is_some() + && let Ok(definitions) = self.find_definition( + handle, + range.start(), + FindPreference { + import_behavior: ImportBehavior::StopAtRenamedImports, + resolve_call_dunders: false, + ..Default::default() + }, + ) + && definitions.iter().any(|definition| { + definition.module.path().as_path() == module.path().as_path() + && matches_definition( + definition.definition_range, + definition.module.code_at(definition.definition_range), + ) + }) + { + references.push(range); + } + } + } references.sort_by_key(|range| range.start()); references.dedup(); Some(references) @@ -6038,6 +6079,23 @@ fn find_global_references_from_definition_impl( // 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)?; + } else if matches!( + definition_kind, + DefinitionMetadata::Attribute | DefinitionMetadata::VariableOrAttribute(_) + ) { + // Slot strings require the candidate's class bindings for attribute resolution. + let slots_handles = candidate_handles + .iter() + .filter(|handle| { + transaction + .bindings_metadata(handle) + .is_some_and(|metadata| metadata.has_class_field(&dunder::SLOTS)) + }) + .cloned() + .collect::>(); + if !slots_handles.is_empty() { + transaction.run_for_handles(&slots_handles, Require::Everything)?; + } } results.extend(process_candidate_handles_with_definition_impl( transaction, From 4f9121ec76b91a424fbf21233b4c3ce935055bd8 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sat, 14 Nov 2026 13:17:20 +0900 Subject: [PATCH 19/19] Index slots occurrences through semantic attribute resolution --- pyrefly/lib/alt/answers.rs | 28 ++++++++ pyrefly/lib/binding/metadata.rs | 6 -- pyrefly/lib/state/lsp.rs | 66 +++---------------- pyrefly/lib/test/lsp/local_find_refs.rs | 6 +- .../test/lsp/lsp_interaction/references.rs | 42 ++++++++++++ 5 files changed, 84 insertions(+), 64 deletions(-) diff --git a/pyrefly/lib/alt/answers.rs b/pyrefly/lib/alt/answers.rs index 4b9534fa3c..36eecbb7e5 100644 --- a/pyrefly/lib/alt/answers.rs +++ b/pyrefly/lib/alt/answers.rs @@ -111,6 +111,8 @@ pub struct Index { /// The call-site range spells the class name, not the definition's own name, so these are /// gated behind `ReferenceOptions::include_constructor_call_sites`. pub constructor_references: SmallMap>, + /// Literal slot occurrences, including attributes declared in this module. + pub slot_references: SmallMap>, /// A map from (child method range) to a list of parent method definitions (ModulePath, parent method). /// This is used to find reimplementations when doing find-references on parent methods. pub parent_methods_map: SmallMap>, @@ -144,6 +146,8 @@ pub enum AttributeReferenceKind { /// constructor protocol, as in `Foo()` reaching `Foo.__init__`. Only class construction /// is recorded this way; calling an instance through its `__call__` is not. ConstructorCall, + /// A literal name in the owning class's `__slots__` binding. + Slot, } /// Visiting a trace store reaches every type recorded in it. @@ -1525,6 +1529,18 @@ impl Answers { } // `pre_solve` has published every exported slot. From this point on, // the preallocated solutions table represents a complete result set. + if self.index.is_some() { + for entry in bindings.dunder_slots_entries() { + if let Some(class) = answers_solver.get_idx(entry.class_idx).0.as_ref() { + answers_solver.record_attribute_definition_index( + &answers_solver.instantiate(class), + &entry.name.id, + entry.name.range, + AttributeReferenceKind::Slot, + ); + } + } + } if let Some(index) = &self.index { let mut index = index.lock(); // Retain export occurrences so unopened reverse dependencies need no AST or answers. @@ -1923,6 +1939,18 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { range: _, docstring_range: _, } => match reference_kind { + AttributeReferenceKind::Slot => index + .lock() + .slot_references + .entry(cls.module_path().dupe()) + .or_default() + .push(( + IndexedAttribute { + cls: cls.index(), + name, + }, + attribute_reference_range, + )), AttributeReferenceKind::ConstructorCall => index .lock() .constructor_references diff --git a/pyrefly/lib/binding/metadata.rs b/pyrefly/lib/binding/metadata.rs index c31396c038..d1254a9a0d 100644 --- a/pyrefly/lib/binding/metadata.rs +++ b/pyrefly/lib/binding/metadata.rs @@ -13,7 +13,6 @@ use pyrefly_types::class::ClassDefIndex; use pyrefly_types::class::ClassFields; -use ruff_python_ast::name::Name; /// Metadata for a single class definition, populated during binding. #[derive(Debug, Clone, Default)] @@ -68,9 +67,4 @@ impl BindingsMetadata { pub fn get_class_mut(&mut self, idx: ClassDefIndex) -> &mut ClassMetadata { &mut self.classes[idx.0 as usize] } - - /// Return whether any class in this module defines a field with this name. - pub fn has_class_field(&self, name: &Name) -> bool { - self.classes.iter().any(|class| class.fields.contains(name)) - } } diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 2341460b2c..569e6e72bb 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -2034,7 +2034,7 @@ impl<'a> Transaction<'a> { } /// The references in `references_by_module` that point at the definition at `definition_range` - /// in `module`. Both `Index` attribute reference maps are keyed and shaped alike, so they share + /// in `module`. The `Index` attribute reference maps are keyed and shaped alike, so they share /// this scan. fn recorded_attribute_references( &self, @@ -5000,45 +5000,16 @@ impl<'a> Transaction<'a> { )); } if is_attribute - && self - .get_bindings_metadata(handle) - .is_some_and(|metadata| metadata.has_class_field(&dunder::SLOTS)) - && let Some(ast) = self.get_ast(handle) + && let Some(index) = self + .get_solutions(handle) + .and_then(|solutions| solutions.get_index()) { - fn collect(expr: &Expr, name: &Name, ranges: &mut Vec) { - if let Expr::StringLiteral(literal) = expr - && literal.value.to_str() == name.as_str() - && let Some(part) = literal.as_single_part_string() - { - ranges.push(part.content_range()); - } - expr.recurse(&mut |expr| collect(expr, name, ranges)); - } - let mut literals = Vec::new(); - ast.visit(&mut |expr| collect(expr, &definition_name, &mut literals)); - let matches_definition = definition_matcher(module, definition_range); - for range in literals { - if dunder_slots_entry(&Ast::locate_node(&ast, range.start())).is_some() - && let Ok(definitions) = self.find_definition( - handle, - range.start(), - FindPreference { - import_behavior: ImportBehavior::StopAtRenamedImports, - resolve_call_dunders: false, - ..Default::default() - }, - ) - && definitions.iter().any(|definition| { - definition.module.path().as_path() == module.path().as_path() - && matches_definition( - definition.definition_range, - definition.module.code_at(definition.definition_range), - ) - }) - { - references.push(range); - } - } + references.extend(self.recorded_attribute_references( + handle, + &index.lock().slot_references, + module, + definition_range, + )); } references.sort_by_key(|range| range.start()); references.dedup(); @@ -6079,23 +6050,6 @@ fn find_global_references_from_definition_impl( // 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)?; - } else if matches!( - definition_kind, - DefinitionMetadata::Attribute | DefinitionMetadata::VariableOrAttribute(_) - ) { - // Slot strings require the candidate's class bindings for attribute resolution. - let slots_handles = candidate_handles - .iter() - .filter(|handle| { - transaction - .bindings_metadata(handle) - .is_some_and(|metadata| metadata.has_class_field(&dunder::SLOTS)) - }) - .cloned() - .collect::>(); - if !slots_handles.is_empty() { - transaction.run_for_handles(&slots_handles, Require::Everything)?; - } } results.extend(process_candidate_handles_with_definition_impl( transaction, diff --git a/pyrefly/lib/test/lsp/local_find_refs.rs b/pyrefly/lib/test/lsp/local_find_refs.rs index 56b3dbd1ce..c5c1f60ead 100644 --- a/pyrefly/lib/test/lsp/local_find_refs.rs +++ b/pyrefly/lib/test/lsp/local_find_refs.rs @@ -296,7 +296,6 @@ fn dunder_slots_references_use_string_contents() { "__slots__ = [r'foo']", "__slots__ = {\"foo\": \"Field documentation.\"}", "__slots__: tuple[str, ...] = (\"\"\"foo\"\"\",)", - "names = __slots__ = (\"foo\",)", ] { let assignment = assignment.replace("foo", name); let code = format!( @@ -439,7 +438,10 @@ fn dunder_slots_global_references_in_indexed_files() { references, vec![("base".to_owned(), 1), ("child".to_owned(), 1)] ); - assert!(transaction.get_answers(&handle("plain")).is_none()); + for name in ["child", "unrelated", "plain"] { + assert!(transaction.get_ast(&handle(name)).is_none()); + assert!(transaction.get_answers(&handle(name)).is_none()); + } } #[test] diff --git a/pyrefly/lib/test/lsp/lsp_interaction/references.rs b/pyrefly/lib/test/lsp/lsp_interaction/references.rs index 413c042030..13d35c36f6 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/references.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/references.rs @@ -139,6 +139,48 @@ fn references_dunder_all_in_unopened_reexports() { interaction.shutdown().unwrap(); } +#[test] +fn references_dunder_slots_in_unopened_subclass() { + let root = TempDir::new().unwrap(); + for (file, code) in [ + ("base.py", "class Base:\n foo: int\n"), + ( + "child.py", + "from base import Base\nclass Child(Base):\n __slots__ = ['foo']\n", + ), + ( + "other.py", + "from base import Base\nclass Other:\n __slots__ = ['foo']\n foo: str\n", + ), + ] { + fs::write(root.path().join(file), code).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().to_path_buf()); + interaction + .initialize(InitializeSettings { + workspace_folders: Some(vec![( + "test".to_owned(), + Uri::from_file_path(root.path()).unwrap(), + )]), + configuration: Some(None), + ..Default::default() + }) + .unwrap(); + interaction.client.did_open("base.py"); + let child_uri = Uri::from_file_path(root.path().join("child.py")).unwrap(); + interaction.client.references("base.py", 1, 5, false).expect_response(json!([ + {"uri": child_uri, "range": {"start": {"line": 2, "character": 18}, "end": {"line": 2, "character": 21}}}, + ])).unwrap(); + interaction.shutdown().unwrap(); +} + // https://github.com/meta/pyrefly/issues/5485 #[test] fn test_references_to_class_member_after_defining_file_shifts() {