diff --git a/pyrefly/lib/lsp/wasm/hover.rs b/pyrefly/lib/lsp/wasm/hover.rs index 01e1f57ae0..7759c15c5a 100644 --- a/pyrefly/lib/lsp/wasm/hover.rs +++ b/pyrefly/lib/lsp/wasm/hover.rs @@ -9,6 +9,7 @@ use std::collections::HashMap; +use dupe::Dupe; use lsp_types::Hover; use lsp_types::HoverContents; use lsp_types::MarkupContent; @@ -315,6 +316,22 @@ fn identifier_text_at( .map(|id| id.identifier.id.to_string()) } +/// Hover should still work when the use site has no type trace, as in +/// `Annotated[..., imported_symbol]` metadata. In that case, recover the type +/// from the resolved definition instead of returning no hover at all. +fn hover_type_from_definition( + transaction: &Transaction<'_>, + current_handle: &Handle, + definition: &FindDefinitionItemWithDocstring, +) -> Option { + let definition_handle = Handle::new( + definition.module.name(), + definition.module.path().dupe(), + current_handle.sys_info().dupe(), + ); + transaction.get_type_at(&definition_handle, definition.definition_range.start()) +} + fn collect_typed_dict_fields_for_hover<'a>( solver: &AnswersSolver>, ty: &Type, @@ -518,8 +535,28 @@ pub fn get_hover( }); } - // Otherwise, fall through to the existing type hover logic - let mut type_ = transaction.get_type_at(handle, position)?; + let definition = transaction + .find_definition( + handle, + position, + FindPreference { + prefer_pyi: false, + ..Default::default() + }, + ) + .map(Vec1::into_vec) + .unwrap_or_default() + .into_iter() + .next(); + + // Otherwise, fall through to the existing type hover logic. Some + // annotation metadata names do not get a type trace at the use site, so + // recover their hover type from the resolved definition. + let mut type_ = transaction.get_type_at(handle, position).or_else(|| { + definition + .as_ref() + .and_then(|definition| hover_type_from_definition(transaction, handle, definition)) + })?; // Helper function to check if we're hovering over a callee and get its range let find_callee_range_at_position = || -> Option { @@ -560,20 +597,7 @@ pub fn get_hover( module, docstring_range, display_name, - }) = transaction - .find_definition( - handle, - position, - FindPreference { - prefer_pyi: false, - ..Default::default() - }, - ) - .map(Vec1::into_vec) - .unwrap_or_default() - // TODO: handle more than 1 definition - .into_iter() - .next() + }) = definition { let kind = metadata.symbol_kind(); let name = { diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 65e4178e35..fd963e5f45 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -1702,13 +1702,32 @@ impl<'a> Transaction<'a> { name: &Name, preference: FindPreference, ) -> Result, EmptyResponseReason> { + let base_type = self.attribute_base_type(handle, base_range)?; + self.find_attribute_definition_for_base_type(handle, preference, base_type, name) + } + + /// Attribute lookup usually relies on a type trace for the base + /// expression, but annotation metadata like `Annotated[..., utils.Name]` + /// may skip tracing the imported module reference. In that case, recover + /// the base type from the exact identifier at `base_range`. + fn attribute_base_type( + &self, + handle: &Handle, + base_range: TextRange, + ) -> Result { let answers = self .get_answers(handle) .ok_or(EmptyResponseReason::AnswersNotFound)?; - let base_type = answers - .get_type_trace(base_range) - .ok_or(EmptyResponseReason::TypeTraceNotFound)?; - self.find_attribute_definition_for_base_type(handle, preference, base_type, name) + if let Some(base_type) = answers.get_type_trace(base_range) { + return Ok(base_type); + } + if let Some(identifier) = self.identifier_at(handle, base_range.start()) + && identifier.identifier.range == base_range + && let Some(base_type) = self.get_type_at(handle, base_range.start()) + { + return Ok(base_type); + } + Err(EmptyResponseReason::TypeTraceNotFound) } pub(crate) fn find_definition_for_imported_module( diff --git a/pyrefly/lib/test/lsp/hover.rs b/pyrefly/lib/test/lsp/hover.rs index 813c78f85d..a80aa20381 100644 --- a/pyrefly/lib/test/lsp/hover.rs +++ b/pyrefly/lib/test/lsp/hover.rs @@ -900,6 +900,56 @@ from lib import bar as baz ); } +#[test] +fn hover_on_qualified_type_alias_in_parameter_annotation() { + let utils = r#" +from typing import Annotated, TypeAlias + +ValueRange: TypeAlias = Annotated[int, "value range"] +"#; + let code = r#" +import utils + +def takes(x: utils.ValueRange) -> None: ... +# ^ +"#; + let report = + get_batched_lsp_operations_report(&[("main", code), ("utils", utils)], get_test_report); + assert!( + !report.contains("\nNone\n"), + "Expected hover for qualified type alias in annotation, got: {report}" + ); + assert!( + report.contains("ValueRange"), + "Expected hover to mention the type alias name, got: {report}" + ); +} + +#[test] +fn hover_on_imported_annotated_metadata_in_parameter_annotation() { + let utils = r#" +class ValueRange: + pass +"#; + let code = r#" +from typing import Annotated +import utils + +def takes(x: Annotated[int, utils.ValueRange]) -> None: ... +# ^ +"#; + let report = + get_batched_lsp_operations_report(&[("main", code), ("utils", utils)], get_test_report); + assert!( + !report.contains("\nNone\n"), + "Expected hover for imported Annotated metadata, got: {report}" + ); + assert!( + report.contains("ValueRange"), + "Expected hover to mention the metadata symbol name, got: {report}" + ); +} + #[test] fn hover_on_import_different_name_alias_second_token_test() { let lib = r#"