Skip to content

Commit 672616c

Browse files
fix
1 parent b430a08 commit 672616c

3 files changed

Lines changed: 113 additions & 20 deletions

File tree

pyrefly/lib/lsp/wasm/hover.rs

Lines changed: 40 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99

1010
use std::collections::HashMap;
1111

12+
use dupe::Dupe;
1213
use lsp_types::Hover;
1314
use lsp_types::HoverContents;
1415
use lsp_types::MarkupContent;
@@ -310,6 +311,22 @@ fn identifier_text_at(
310311
.map(|id| id.identifier.id.to_string())
311312
}
312313

314+
/// Hover should still work when the use site has no type trace, as in
315+
/// `Annotated[..., imported_symbol]` metadata. In that case, recover the type
316+
/// from the resolved definition instead of returning no hover at all.
317+
fn hover_type_from_definition(
318+
transaction: &Transaction<'_>,
319+
current_handle: &Handle,
320+
definition: &FindDefinitionItemWithDocstring,
321+
) -> Option<Type> {
322+
let definition_handle = Handle::new(
323+
definition.module.name(),
324+
definition.module.path().dupe(),
325+
current_handle.sys_info().dupe(),
326+
);
327+
transaction.get_type_at(&definition_handle, definition.definition_range.start())
328+
}
329+
313330
fn collect_typed_dict_fields_for_hover<'a>(
314331
solver: &AnswersSolver<TransactionHandle<'a>>,
315332
ty: &Type,
@@ -513,8 +530,28 @@ pub fn get_hover(
513530
});
514531
}
515532

516-
// Otherwise, fall through to the existing type hover logic
517-
let mut type_ = transaction.get_type_at(handle, position)?;
533+
let definition = transaction
534+
.find_definition(
535+
handle,
536+
position,
537+
FindPreference {
538+
prefer_pyi: false,
539+
..Default::default()
540+
},
541+
)
542+
.map(Vec1::into_vec)
543+
.unwrap_or_default()
544+
.into_iter()
545+
.next();
546+
547+
// Otherwise, fall through to the existing type hover logic. Some
548+
// annotation metadata names do not get a type trace at the use site, so
549+
// recover their hover type from the resolved definition.
550+
let mut type_ = transaction.get_type_at(handle, position).or_else(|| {
551+
definition
552+
.as_ref()
553+
.and_then(|definition| hover_type_from_definition(transaction, handle, definition))
554+
})?;
518555

519556
// Helper function to check if we're hovering over a callee and get its range
520557
let find_callee_range_at_position = || -> Option<TextRange> {
@@ -555,20 +592,7 @@ pub fn get_hover(
555592
module,
556593
docstring_range,
557594
display_name,
558-
}) = transaction
559-
.find_definition(
560-
handle,
561-
position,
562-
FindPreference {
563-
prefer_pyi: false,
564-
..Default::default()
565-
},
566-
)
567-
.map(Vec1::into_vec)
568-
.unwrap_or_default()
569-
// TODO: handle more than 1 definition
570-
.into_iter()
571-
.next()
595+
}) = definition
572596
{
573597
let kind = metadata.symbol_kind();
574598
let name = {

pyrefly/lib/state/lsp.rs

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1702,13 +1702,32 @@ impl<'a> Transaction<'a> {
17021702
name: &Name,
17031703
preference: FindPreference,
17041704
) -> Result<Vec1<FindDefinitionItemWithDocstring>, EmptyResponseReason> {
1705+
let base_type = self.attribute_base_type(handle, base_range)?;
1706+
self.find_attribute_definition_for_base_type(handle, preference, base_type, name)
1707+
}
1708+
1709+
/// Attribute lookup usually relies on a type trace for the base
1710+
/// expression, but annotation metadata like `Annotated[..., utils.Name]`
1711+
/// may skip tracing the imported module reference. In that case, recover
1712+
/// the base type from the exact identifier at `base_range`.
1713+
fn attribute_base_type(
1714+
&self,
1715+
handle: &Handle,
1716+
base_range: TextRange,
1717+
) -> Result<Type, EmptyResponseReason> {
17051718
let answers = self
17061719
.get_answers(handle)
17071720
.ok_or(EmptyResponseReason::AnswersNotFound)?;
1708-
let base_type = answers
1709-
.get_type_trace(base_range)
1710-
.ok_or(EmptyResponseReason::TypeTraceNotFound)?;
1711-
self.find_attribute_definition_for_base_type(handle, preference, base_type, name)
1721+
if let Some(base_type) = answers.get_type_trace(base_range) {
1722+
return Ok(base_type);
1723+
}
1724+
if let Some(identifier) = self.identifier_at(handle, base_range.start())
1725+
&& identifier.identifier.range == base_range
1726+
&& let Some(base_type) = self.get_type_at(handle, base_range.start())
1727+
{
1728+
return Ok(base_type);
1729+
}
1730+
Err(EmptyResponseReason::TypeTraceNotFound)
17121731
}
17131732

17141733
pub(crate) fn find_definition_for_imported_module(

pyrefly/lib/test/lsp/hover.rs

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -900,6 +900,56 @@ from lib import bar as baz
900900
);
901901
}
902902

903+
#[test]
904+
fn hover_on_qualified_type_alias_in_parameter_annotation() {
905+
let utils = r#"
906+
from typing import Annotated, TypeAlias
907+
908+
ValueRange: TypeAlias = Annotated[int, "value range"]
909+
"#;
910+
let code = r#"
911+
import utils
912+
913+
def takes(x: utils.ValueRange) -> None: ...
914+
# ^
915+
"#;
916+
let report =
917+
get_batched_lsp_operations_report(&[("main", code), ("utils", utils)], get_test_report);
918+
assert!(
919+
!report.contains("\nNone\n"),
920+
"Expected hover for qualified type alias in annotation, got: {report}"
921+
);
922+
assert!(
923+
report.contains("ValueRange"),
924+
"Expected hover to mention the type alias name, got: {report}"
925+
);
926+
}
927+
928+
#[test]
929+
fn hover_on_imported_annotated_metadata_in_parameter_annotation() {
930+
let utils = r#"
931+
class ValueRange:
932+
pass
933+
"#;
934+
let code = r#"
935+
from typing import Annotated
936+
import utils
937+
938+
def takes(x: Annotated[int, utils.ValueRange]) -> None: ...
939+
# ^
940+
"#;
941+
let report =
942+
get_batched_lsp_operations_report(&[("main", code), ("utils", utils)], get_test_report);
943+
assert!(
944+
!report.contains("\nNone\n"),
945+
"Expected hover for imported Annotated metadata, got: {report}"
946+
);
947+
assert!(
948+
report.contains("ValueRange"),
949+
"Expected hover to mention the metadata symbol name, got: {report}"
950+
);
951+
}
952+
903953
#[test]
904954
fn hover_on_import_different_name_alias_second_token_test() {
905955
let lib = r#"

0 commit comments

Comments
 (0)