fix doesn't understand usage of inline Annotated function parameters #2531

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<Type> {
+    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<TransactionHandle<'a>>,
     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<TextRange> {
@@ -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 = {

pyrefly/lib/state/lsp.rs

@@ -1702,13 +1702,32 @@ impl<'a> Transaction<'a> {
         name: &Name,
         preference: FindPreference,
     ) -> Result<Vec1<FindDefinitionItemWithDocstring>, 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<Type, EmptyResponseReason> {
         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(

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#"