diff --git a/crates/pyrefly_types/src/display.rs b/crates/pyrefly_types/src/display.rs index 8b66c13a4..b9f2e64ad 100644 --- a/crates/pyrefly_types/src/display.rs +++ b/crates/pyrefly_types/src/display.rs @@ -7,10 +7,10 @@ //! Display a type. The complexity comes from if we have two classes with the same name, //! we want to display disambiguating information (e.g. module name or location). +use std::borrow::Cow; use std::fmt; use std::fmt::Display; -use pyrefly_python::module::TextRangeWithModule; use pyrefly_python::module_name::ModuleName; use pyrefly_python::qname::QName; use pyrefly_util::display::Fmt; @@ -29,6 +29,7 @@ use crate::literal::Lit; use crate::tuple::Tuple; use crate::type_output::DisplayOutput; use crate::type_output::OutputWithLocations; +use crate::type_output::TypeLabelPart; use crate::type_output::TypeOutput; use crate::types::AnyStyle; use crate::types::BoundMethod; @@ -237,11 +238,12 @@ impl<'a> TypeDisplayContext<'a> { name: &str, output: &mut impl TypeOutput, ) -> fmt::Result { - if self.always_display_module_name { - output.write_str(&format!("{}.{}", module, name)) - } else { - output.write_str(name) - } + let module_name = ModuleName::from_str(module); + output.write_symbol( + module_name, + Cow::Borrowed(name), + self.always_display_module_name, + ) } /// Helper function to format a sequence of types with a separator. @@ -695,7 +697,10 @@ impl<'a> TypeDisplayContext<'a> { output.write_str(&format!("{q}"))?; output.write_str("]") } - Type::SpecialForm(x) => output.write_str(&format!("{x}")), + Type::SpecialForm(x) => { + let name = x.to_string(); + self.maybe_fmt_with_module("typing", &name, output) + } Type::Ellipsis => output.write_str("Ellipsis"), Type::Any(style) => match style { AnyStyle::Explicit => self.maybe_fmt_with_module("typing", "Any", output), @@ -781,7 +786,7 @@ impl Type { c.display(self).to_string() } - pub fn get_types_with_locations(&self) -> Vec<(String, Option)> { + pub fn get_types_with_locations(&self) -> Vec { let ctx = TypeDisplayContext::new(&[self]); let mut output = OutputWithLocations::new(&ctx); ctx.fmt_helper_generic(self, false, &mut output).unwrap(); @@ -1651,32 +1656,30 @@ def overloaded_func[T]( } // Helper functions for testing get_types_with_location - fn get_parts(t: &Type) -> Vec<(String, Option)> { + fn get_parts(t: &Type) -> Vec { let ctx = TypeDisplayContext::new(&[t]); let output = ctx.get_types_with_location(t, false); output.parts().to_vec() } - fn parts_to_string(parts: &[(String, Option)]) -> String { - parts.iter().map(|(s, _)| s.as_str()).collect::() + fn parts_to_string(parts: &[TypeLabelPart]) -> String { + parts + .iter() + .map(|part| part.text.as_str()) + .collect::() } - fn assert_part_has_location( - parts: &[(String, Option)], - name: &str, - module: &str, - position: u32, - ) { - let part = parts.iter().find(|(s, _)| s == name); + fn assert_part_has_location(parts: &[TypeLabelPart], name: &str, module: &str, position: u32) { + let part = parts.iter().find(|part| part.text == name); assert!(part.is_some(), "Should have {} in parts", name); - let (_, location) = part.unwrap(); - assert!(location.is_some(), "{} should have location", name); - let loc = location.as_ref().unwrap(); + let loc = part.unwrap().location.as_ref(); + assert!(loc.is_some(), "{} should have location", name); + let loc = loc.unwrap(); assert_eq!(loc.module.name().as_str(), module); assert_eq!(loc.range.start().to_u32(), position); } - fn assert_output_contains(parts: &[(String, Option)], needle: &str) { + fn assert_output_contains(parts: &[TypeLabelPart], needle: &str) { let full_str = parts_to_string(parts); assert!( full_str.contains(needle), @@ -1705,9 +1708,12 @@ def overloaded_func[T]( let t = Type::ClassType(ClassType::new(foo, TArgs::new(tparams, vec![inner_type]))); let parts = get_parts(&t); - assert_eq!(parts[0].0, "Foo"); + assert_eq!(parts[0].text, "Foo"); assert_part_has_location(&parts, "Foo", "test.module", 10); - assert!(parts.iter().any(|(s, _)| s == "Bar"), "Should have Bar"); + assert!( + parts.iter().any(|part| part.text == "Bar"), + "Should have Bar" + ); } #[test] @@ -1716,8 +1722,11 @@ def overloaded_func[T]( let t = tvar.to_type(); let parts = get_parts(&t); - assert_eq!(parts[0].0, "TypeVar["); - assert!(parts[0].1.is_none(), "TypeVar[ should not have location"); + assert_eq!(parts[0].text, "TypeVar["); + assert!( + parts[0].location.is_none(), + "TypeVar[ should not have location" + ); assert_part_has_location(&parts, "T", "test.module", 15); } @@ -1733,8 +1742,14 @@ def overloaded_func[T]( let parts1 = ctx.get_types_with_location(&t1, false).parts().to_vec(); let parts2 = ctx.get_types_with_location(&t2, false).parts().to_vec(); - let loc1 = parts1.iter().find_map(|(_, loc)| loc.as_ref()).unwrap(); - let loc2 = parts2.iter().find_map(|(_, loc)| loc.as_ref()).unwrap(); + let loc1 = parts1 + .iter() + .find_map(|part| part.location.as_ref()) + .unwrap(); + let loc2 = parts2 + .iter() + .find_map(|part| part.location.as_ref()) + .unwrap(); assert_ne!( loc1.range.start().to_u32(), loc2.range.start().to_u32(), @@ -1749,6 +1764,11 @@ def overloaded_func[T]( assert_output_contains(&parts, "Literal"); assert_output_contains(&parts, "True"); + let literal_part = parts.iter().find(|part| part.text == "Literal"); + assert!(literal_part.is_some()); + let symbol = literal_part.unwrap().symbol.as_ref().unwrap(); + assert_eq!(symbol.module.as_str(), "typing"); + assert_eq!(symbol.name, "Literal"); } #[test] @@ -1773,8 +1793,8 @@ def overloaded_func[T]( let parts = get_parts(&t); assert_eq!(parts.len(), 1); - assert_eq!(parts[0].0, "None"); - assert!(parts[0].1.is_none(), "None should not have location"); + assert_eq!(parts[0].text, "None"); + assert!(parts[0].location.is_none(), "None should not have location"); } #[test] @@ -1800,7 +1820,11 @@ def overloaded_func[T]( for param in &["T", "U", "Ts"] { assert_output_contains(&parts, param); } - assert!(parts.iter().any(|(s, loc)| s == "[" && loc.is_none())); + assert!( + parts + .iter() + .any(|part| part.text == "[" && part.location.is_none()) + ); assert!(parts_to_string(&parts).starts_with('[')); assert_output_contains(&parts, "]("); } @@ -1830,7 +1854,7 @@ def overloaded_func[T]( for expected in &["Literal", "Color", "RED"] { assert_output_contains(&parts, expected); } - assert!(parts.iter().any(|(_, loc)| loc.is_some())); + assert!(parts.iter().any(|part| part.location.is_some())); } #[test] diff --git a/crates/pyrefly_types/src/type_output.rs b/crates/pyrefly_types/src/type_output.rs index c3ef11fe6..cd20e37ae 100644 --- a/crates/pyrefly_types/src/type_output.rs +++ b/crates/pyrefly_types/src/type_output.rs @@ -5,9 +5,11 @@ * LICENSE file in the root directory of this source tree. */ +use std::borrow::Cow; use std::fmt; use pyrefly_python::module::TextRangeWithModule; +use pyrefly_python::module_name::ModuleName; use pyrefly_python::qname::QName; use crate::display::TypeDisplayContext; @@ -26,6 +28,12 @@ pub trait TypeOutput { fn write_lit(&mut self, lit: &Lit) -> fmt::Result; fn write_targs(&mut self, targs: &TArgs) -> fmt::Result; fn write_type(&mut self, ty: &Type) -> fmt::Result; + fn write_symbol( + &mut self, + module: ModuleName, + name: Cow<'_, str>, + display_module: bool, + ) -> fmt::Result; } /// Implementation of `TypeOutput` that writes formatted types to plain text. @@ -63,6 +71,22 @@ impl<'a, 'b, 'f> TypeOutput for DisplayOutput<'a, 'b, 'f> { fn write_type(&mut self, ty: &Type) -> fmt::Result { write!(self.formatter, "{}", self.context.display(ty)) } + + fn write_symbol( + &mut self, + module: ModuleName, + name: Cow<'_, str>, + display_module: bool, + ) -> fmt::Result { + if display_module { + write!(self.formatter, "{}.{}", module, name) + } else { + match name { + Cow::Borrowed(s) => self.formatter.write_str(s), + Cow::Owned(s) => self.formatter.write_str(&s), + } + } + } } /// This struct is used to collect the type to be displayed as a vector. Each element @@ -71,8 +95,21 @@ impl<'a, 'b, 'f> TypeOutput for DisplayOutput<'a, 'b, 'f> { /// each of these will be concatenated to create the final type. /// The second element of the vector is an optional location. For any part that do have /// a location this will be included. For separators like '|', '[', etc. this will be None. +#[derive(Clone, Debug)] +pub struct TypeSymbolReference { + pub module: ModuleName, + pub name: String, +} + +#[derive(Clone, Debug)] +pub struct TypeLabelPart { + pub text: String, + pub location: Option, + pub symbol: Option, +} + pub struct OutputWithLocations<'a> { - parts: Vec<(String, Option)>, + parts: Vec, context: &'a TypeDisplayContext<'a>, } @@ -84,20 +121,28 @@ impl<'a> OutputWithLocations<'a> { } } - pub fn parts(&self) -> &[(String, Option)] { + pub fn parts(&self) -> &[TypeLabelPart] { &self.parts } } impl TypeOutput for OutputWithLocations<'_> { fn write_str(&mut self, s: &str) -> fmt::Result { - self.parts.push((s.to_owned(), None)); + self.parts.push(TypeLabelPart { + text: s.to_owned(), + location: None, + symbol: None, + }); Ok(()) } fn write_qname(&mut self, qname: &QName) -> fmt::Result { let location = TextRangeWithModule::new(qname.module().clone(), qname.range()); - self.parts.push((qname.id().to_string(), Some(location))); + self.parts.push(TypeLabelPart { + text: qname.id().to_string(), + location: Some(location), + symbol: None, + }); Ok(()) } @@ -115,7 +160,11 @@ impl TypeOutput for OutputWithLocations<'_> { } _ => None, }; - self.parts.push((formatted, location)); + self.parts.push(TypeLabelPart { + text: formatted, + location, + symbol: None, + }); Ok(()) } @@ -140,6 +189,31 @@ impl TypeOutput for OutputWithLocations<'_> { // Format the type and extract location if it has a qname self.context.fmt_helper_generic(ty, false, self) } + + fn write_symbol( + &mut self, + module: ModuleName, + name: Cow<'_, str>, + display_module: bool, + ) -> fmt::Result { + let text = if display_module { + format!("{}.{}", module, name) + } else { + match &name { + Cow::Borrowed(s) => (*s).to_owned(), + Cow::Owned(s) => s.clone(), + } + }; + self.parts.push(TypeLabelPart { + text, + location: None, + symbol: Some(TypeSymbolReference { + module, + name: name.into_owned(), + }), + }); + Ok(()) + } } #[cfg(test)] @@ -198,17 +272,17 @@ mod tests { output.write_str("hello").unwrap(); assert_eq!(output.parts().len(), 1); - assert_eq!(output.parts()[0].0, "hello"); - assert!(output.parts()[0].1.is_none()); + assert_eq!(output.parts()[0].text, "hello"); + assert!(output.parts()[0].location.is_none()); output.write_str(" world").unwrap(); assert_eq!(output.parts().len(), 2); - assert_eq!(output.parts()[1].0, " world"); - assert!(output.parts()[1].1.is_none()); + assert_eq!(output.parts()[1].text, " world"); + assert!(output.parts()[1].location.is_none()); let parts = output.parts(); - assert_eq!(parts[0].0, "hello"); - assert_eq!(parts[1].0, " world"); + assert_eq!(parts[0].text, "hello"); + assert_eq!(parts[1].text, " world"); } #[test] @@ -231,11 +305,11 @@ mod tests { output.write_qname(&qname).unwrap(); assert_eq!(output.parts().len(), 1); - let (name_str, location) = &output.parts()[0]; - assert_eq!(name_str, "MyClass"); + let part = &output.parts()[0]; + assert_eq!(part.text, "MyClass"); - assert!(location.is_some()); - let loc = location.as_ref().unwrap(); + assert!(part.location.is_some()); + let loc = part.location.as_ref().unwrap(); assert_eq!( loc.range, TextRange::new(TextSize::new(4), TextSize::new(11)) @@ -253,8 +327,8 @@ mod tests { output.write_lit(&str_lit).unwrap(); assert_eq!(output.parts().len(), 1); - assert_eq!(output.parts()[0].0, "'hello'"); - assert!(output.parts()[0].1.is_none()); + assert_eq!(output.parts()[0].text, "'hello'"); + assert!(output.parts()[0].location.is_none()); } #[test] @@ -274,13 +348,13 @@ mod tests { output.write_lit(&enum_lit).unwrap(); assert_eq!(output.parts().len(), 1); - let (formatted, location) = &output.parts()[0]; + let part = &output.parts()[0]; - assert_eq!(formatted, "Color.RED"); + assert_eq!(part.text, "Color.RED"); // Verify the location was captured from the enum's class qname - assert!(location.is_some()); - let loc = location.as_ref().unwrap(); + assert!(part.location.is_some()); + let loc = part.location.as_ref().unwrap(); assert_eq!(loc.range, TextRange::empty(TextSize::new(10))); assert_eq!(loc.module.name(), ModuleName::from_str("colors")); } @@ -333,26 +407,26 @@ mod tests { // Now that write_type is implemented, it actually writes the types // Should have: "[", "None", ", ", "LiteralString", ", ", "Any", "]" assert_eq!(output.parts().len(), 7); - assert_eq!(output.parts()[0].0, "["); - assert!(output.parts()[0].1.is_none()); + assert_eq!(output.parts()[0].text, "["); + assert!(output.parts()[0].location.is_none()); - assert_eq!(output.parts()[1].0, "None"); - assert!(output.parts()[1].1.is_none()); + assert_eq!(output.parts()[1].text, "None"); + assert!(output.parts()[1].location.is_none()); - assert_eq!(output.parts()[2].0, ", "); - assert!(output.parts()[2].1.is_none()); + assert_eq!(output.parts()[2].text, ", "); + assert!(output.parts()[2].location.is_none()); - assert_eq!(output.parts()[3].0, "LiteralString"); - assert!(output.parts()[3].1.is_none()); + assert_eq!(output.parts()[3].text, "LiteralString"); + assert!(output.parts()[3].location.is_none()); - assert_eq!(output.parts()[4].0, ", "); - assert!(output.parts()[4].1.is_none()); + assert_eq!(output.parts()[4].text, ", "); + assert!(output.parts()[4].location.is_none()); - assert_eq!(output.parts()[5].0, "Any"); - assert!(output.parts()[5].1.is_none()); + assert_eq!(output.parts()[5].text, "Any"); + assert!(output.parts()[5].location.is_none()); - assert_eq!(output.parts()[6].0, "]"); - assert!(output.parts()[6].1.is_none()); + assert_eq!(output.parts()[6].text, "]"); + assert!(output.parts()[6].location.is_none()); } #[test] @@ -363,13 +437,13 @@ mod tests { // Test simple types that don't have locations output.write_type(&Type::None).unwrap(); assert_eq!(output.parts().len(), 1); - assert_eq!(output.parts()[0].0, "None"); - assert!(output.parts()[0].1.is_none()); + assert_eq!(output.parts()[0].text, "None"); + assert!(output.parts()[0].location.is_none()); output.write_type(&Type::LiteralString).unwrap(); assert_eq!(output.parts().len(), 2); - assert_eq!(output.parts()[1].0, "LiteralString"); - assert!(output.parts()[1].1.is_none()); + assert_eq!(output.parts()[1].text, "LiteralString"); + assert!(output.parts()[1].location.is_none()); } #[test] @@ -389,13 +463,13 @@ mod tests { assert!(!output.parts().is_empty()); // Find the int and str parts and verify they have locations - let int_part = output.parts().iter().find(|p| p.0 == "int"); + let int_part = output.parts().iter().find(|p| p.text == "int"); assert!(int_part.is_some()); - assert!(int_part.unwrap().1.is_some()); + assert!(int_part.unwrap().location.is_some()); - let str_part = output.parts().iter().find(|p| p.0 == "str"); + let str_part = output.parts().iter().find(|p| p.text == "str"); assert!(str_part.is_some()); - assert!(str_part.unwrap().1.is_some()); + assert!(str_part.unwrap().location.is_some()); } #[test] @@ -414,7 +488,11 @@ mod tests { ctx.fmt_helper_generic(&union_type, false, &mut output) .unwrap(); - let parts_str: String = output.parts().iter().map(|(s, _)| s.as_str()).collect(); + let parts_str: String = output + .parts() + .iter() + .map(|part| part.text.as_str()) + .collect(); assert_eq!(parts_str, "int | str | None"); // New behavior: Union types are split into separate parts @@ -423,20 +501,26 @@ mod tests { assert_eq!(parts.len(), 5, "Union should be split into 5 parts"); // Verify each part - assert_eq!(parts[0].0, "int"); - assert!(parts[0].1.is_some(), "int should have location"); + assert_eq!(parts[0].text, "int"); + assert!(parts[0].location.is_some(), "int should have location"); - assert_eq!(parts[1].0, " | "); - assert!(parts[1].1.is_none(), "separator should not have location"); + assert_eq!(parts[1].text, " | "); + assert!( + parts[1].location.is_none(), + "separator should not have location" + ); - assert_eq!(parts[2].0, "str"); - assert!(parts[2].1.is_some(), "str should have location"); + assert_eq!(parts[2].text, "str"); + assert!(parts[2].location.is_some(), "str should have location"); - assert_eq!(parts[3].0, " | "); - assert!(parts[3].1.is_none(), "separator should not have location"); + assert_eq!(parts[3].text, " | "); + assert!( + parts[3].location.is_none(), + "separator should not have location" + ); - assert_eq!(parts[4].0, "None"); - assert!(parts[4].1.is_none(), "None should not have location"); + assert_eq!(parts[4].text, "None"); + assert!(parts[4].location.is_none(), "None should not have location"); } #[test] @@ -461,7 +545,11 @@ mod tests { .unwrap(); // Check the concatenated result - let parts_str: String = output.parts().iter().map(|(s, _)| s.as_str()).collect(); + let parts_str: String = output + .parts() + .iter() + .map(|part| part.text.as_str()) + .collect(); assert_eq!(parts_str, "int & str"); // New behavior: Intersection types are split into separate parts @@ -470,14 +558,17 @@ mod tests { assert_eq!(parts.len(), 3, "Intersection should be split into 3 parts"); // Verify each part - assert_eq!(parts[0].0, "int"); - assert!(parts[0].1.is_some(), "int should have location"); + assert_eq!(parts[0].text, "int"); + assert!(parts[0].location.is_some(), "int should have location"); - assert_eq!(parts[1].0, " & "); - assert!(parts[1].1.is_none(), "separator should not have location"); + assert_eq!(parts[1].text, " & "); + assert!( + parts[1].location.is_none(), + "separator should not have location" + ); - assert_eq!(parts[2].0, "str"); - assert!(parts[2].1.is_some(), "str should have location"); + assert_eq!(parts[2].text, "str"); + assert!(parts[2].location.is_some(), "str should have location"); } #[test] @@ -497,7 +588,11 @@ mod tests { ctx.fmt_helper_generic(&tuple_type, false, &mut output) .unwrap(); - let parts_str: String = output.parts().iter().map(|(s, _)| s.as_str()).collect(); + let parts_str: String = output + .parts() + .iter() + .map(|part| part.text.as_str()) + .collect(); assert_eq!(parts_str, "tuple[int]"); // Current behavior: The "tuple" part is NOT clickable @@ -506,20 +601,23 @@ mod tests { assert_eq!(parts.len(), 4, "Should have 4 parts"); // Verify each part - assert_eq!(parts[0].0, "tuple"); + assert_eq!(parts[0].text, "tuple"); assert!( - parts[0].1.is_none(), + parts[0].location.is_none(), "tuple[ should not have location (not clickable)" ); - assert_eq!(parts[1].0, "["); - assert!(parts[1].1.is_none(), "[ should not have location"); + assert_eq!(parts[1].text, "["); + assert!(parts[1].location.is_none(), "[ should not have location"); - assert_eq!(parts[2].0, "int"); - assert!(parts[2].1.is_some(), "int should have location (clickable)"); + assert_eq!(parts[2].text, "int"); + assert!( + parts[2].location.is_some(), + "int should have location (clickable)" + ); - assert_eq!(parts[3].0, "]"); - assert!(parts[3].1.is_none(), "] should not have location"); + assert_eq!(parts[3].text, "]"); + assert!(parts[3].location.is_none(), "] should not have location"); } #[test] @@ -537,7 +635,11 @@ mod tests { ctx.fmt_helper_generic(&literal_type, false, &mut output) .unwrap(); - let parts_str: String = output.parts().iter().map(|(s, _)| s.as_str()).collect(); + let parts_str: String = output + .parts() + .iter() + .map(|part| part.text.as_str()) + .collect(); assert_eq!(parts_str, "Literal[1]"); // Current behavior: The "Literal" part is NOT clickable @@ -546,19 +648,19 @@ mod tests { assert_eq!(parts.len(), 4, "Should have 4 parts"); // Verify each part - assert_eq!(parts[0].0, "Literal"); + assert_eq!(parts[0].text, "Literal"); assert!( - parts[0].1.is_none(), + parts[0].location.is_none(), "Literal should not have location (not clickable)" ); - assert_eq!(parts[1].0, "["); - assert!(parts[1].1.is_none(), "[ should not have location"); + assert_eq!(parts[1].text, "["); + assert!(parts[1].location.is_none(), "[ should not have location"); - assert_eq!(parts[2].0, "1"); - assert!(parts[2].1.is_none(), "1 should not have location"); + assert_eq!(parts[2].text, "1"); + assert!(parts[2].location.is_none(), "1 should not have location"); - assert_eq!(parts[3].0, "]"); - assert!(parts[3].1.is_none(), "] should not have location"); + assert_eq!(parts[3].text, "]"); + assert!(parts[3].location.is_none(), "] should not have location"); } } diff --git a/pyrefly/lib/lsp/wasm/inlay_hints.rs b/pyrefly/lib/lsp/wasm/inlay_hints.rs index 56766aa56..fee876167 100644 --- a/pyrefly/lib/lsp/wasm/inlay_hints.rs +++ b/pyrefly/lib/lsp/wasm/inlay_hints.rs @@ -13,6 +13,7 @@ use pyrefly_python::ast::Ast; use pyrefly_python::module::TextRangeWithModule; use pyrefly_types::literal::Lit; use pyrefly_types::literal::LitEnum; +use pyrefly_types::type_output::TypeSymbolReference; use pyrefly_util::visit::Visit; use ruff_python_ast::Expr; use ruff_python_ast::ExprAttribute; @@ -28,6 +29,7 @@ use ruff_text_size::TextSize; use crate::binding::binding::Binding; use crate::binding::binding::Key; +use crate::export::exports::ExportLocation; use crate::state::lsp::AllOffPartial; use crate::state::lsp::InlayHintConfig; use crate::state::state::CancellableTransaction; @@ -76,6 +78,38 @@ pub fn normalize_singleton_function_type_into_params(type_: Type) -> Option Transaction<'a> { + fn resolve_type_symbol_reference( + &self, + handle: &Handle, + symbol: &TypeSymbolReference, + ) -> Option { + const MAX_LOOKUP_DEPTH: usize = 8; + let mut current_handle = handle.clone(); + let mut current_module = symbol.module; + let mut current_name = ruff_python_ast::name::Name::new(symbol.name.clone()); + for _ in 0..MAX_LOOKUP_DEPTH { + let module_handle = self + .import_handle(¤t_handle, current_module, None) + .finding()?; + let module_info = self.get_module_info(&module_handle)?; + let exports = self.get_exports(&module_handle); + match exports.get(¤t_name) { + Some(ExportLocation::ThisModule(export)) => { + return Some(TextRangeWithModule::new(module_info, export.location)); + } + Some(ExportLocation::OtherModule(next_module, alias)) => { + current_handle = module_handle; + current_module = *next_module; + if let Some(alias_name) = alias { + current_name = alias_name.clone(); + } + } + None => return None, + } + } + None + } + pub fn inlay_hints( &self, handle: &Handle, @@ -142,11 +176,16 @@ impl<'a> Transaction<'a> { // Use get_types_with_locations to get type parts with location info let type_parts = ty.get_types_with_locations(); let label_parts = once((" -> ".to_owned(), None)) - .chain( - type_parts - .iter() - .map(|(text, loc)| (text.clone(), loc.clone())), - ) + .chain(type_parts.iter().map(|part| { + let location = part.location.clone().or_else(|| { + part.symbol.as_ref().and_then(|symbol| { + self.resolve_type_symbol_reference( + handle, symbol, + ) + }) + }); + (part.text.clone(), location) + })) .collect(); res.push((fun.def.parameters.range.end(), label_parts)); } @@ -184,11 +223,14 @@ impl<'a> Transaction<'a> { // Use get_types_with_locations to get type parts with location info let type_parts = ty.get_types_with_locations(); let label_parts = once((": ".to_owned(), None)) - .chain( - type_parts - .iter() - .map(|(text, loc)| (text.clone(), loc.clone())), - ) + .chain(type_parts.iter().map(|part| { + let location = part.location.clone().or_else(|| { + part.symbol.as_ref().and_then(|symbol| { + self.resolve_type_symbol_reference(handle, symbol) + }) + }); + (part.text.clone(), location) + })) .collect(); res.push((key.range().end(), label_parts)); } diff --git a/pyrefly/lib/test/lsp/lsp_interaction/inlay_hint.rs b/pyrefly/lib/test/lsp/lsp_interaction/inlay_hint.rs index d3ee47864..e8f06f70d 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/inlay_hint.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/inlay_hint.rs @@ -9,7 +9,10 @@ use serde_json::json; use crate::test::lsp::lsp_interaction::object_model::InitializeSettings; use crate::test::lsp::lsp_interaction::object_model::LspInteraction; +use crate::test::lsp::lsp_interaction::util::ExpectedInlayHint; +use crate::test::lsp::lsp_interaction::util::ExpectedTextEdit; use crate::test::lsp::lsp_interaction::util::get_test_files_root; +use crate::test::lsp::lsp_interaction::util::inlay_hints_match_expected; #[test] fn test_inlay_hint_default_config() { @@ -25,69 +28,44 @@ fn test_inlay_hint_default_config() { interaction.client.did_open("inlay_hint_test.py"); + let expected = [ + ExpectedInlayHint { + labels: &[ + " -> ", "tuple", "[", "Literal", "[", "1", "]", ", ", "Literal", "[", "2", "]", "]", + ], + position: (6, 21), + text_edit: ExpectedTextEdit { + new_text: " -> tuple[Literal[1], Literal[2]]", + range_start: (6, 21), + range_end: (6, 21), + }, + }, + ExpectedInlayHint { + labels: &[ + ": ", "tuple", "[", "Literal", "[", "1", "]", ", ", "Literal", "[", "2", "]", "]", + ], + position: (11, 6), + text_edit: ExpectedTextEdit { + new_text: ": tuple[Literal[1], Literal[2]]", + range_start: (11, 6), + range_end: (11, 6), + }, + }, + ExpectedInlayHint { + labels: &[" -> ", "Literal", "[", "0", "]"], + position: (14, 15), + text_edit: ExpectedTextEdit { + new_text: " -> Literal[0]", + range_start: (14, 15), + range_end: (14, 15), + }, + }, + ]; + interaction .client .inlay_hint("inlay_hint_test.py", 0, 0, 100, 0) - .expect_response(json!([ - { - "label":[ - {"value":" -> "}, - {"value":"tuple"}, - {"value":"["}, - {"value":"Literal"}, - {"value":"["}, - {"value":"1"}, - {"value":"]"}, - {"value":", "}, - {"value":"Literal"}, - {"value":"["}, - {"value":"2"}, - {"value":"]"}, - {"value":"]"} - ], - "position":{"character":21,"line":6}, - "textEdits":[{ - "newText":" -> tuple[Literal[1], Literal[2]]", - "range":{"end":{"character":21,"line":6},"start":{"character":21,"line":6}} - }] - }, - { - "label":[ - {"value":": "}, - {"value":"tuple"}, - {"value":"["}, - {"value":"Literal"}, - {"value":"["}, - {"value":"1"}, - {"value":"]"}, - {"value":", "}, - {"value":"Literal"}, - {"value":"["}, - {"value":"2"}, - {"value":"]"}, - {"value":"]"} - ], - "position":{"character":6,"line":11}, - "textEdits":[{ - "newText":": tuple[Literal[1], Literal[2]]", - "range":{"end":{"character":6,"line":11},"start":{"character":6,"line":11}} - }] - }, - { - "label":[ - {"value":" -> "}, - {"value":"Literal"}, - {"value":"["}, - {"value":"0"}, - {"value":"]"} - ], - "position":{"character":15,"line":14}, - "textEdits":[{ - "newText":" -> Literal[0]", - "range":{"end":{"character":15,"line":14},"start":{"character":15,"line":14}} - }] - } - ])) + .expect_response_with(|result| inlay_hints_match_expected(result, &expected)) .unwrap(); interaction.shutdown().unwrap(); @@ -181,42 +159,32 @@ fn test_inlay_hint_disable_variables() { interaction .client .inlay_hint("inlay_hint_test.py", 0, 0, 100, 0) - .expect_response(json!([{ - "label":[ - {"value":" -> "}, - {"value":"tuple"}, - {"value":"["}, - {"value":"Literal"}, - {"value":"["}, - {"value":"1"}, - {"value":"]"}, - {"value":", "}, - {"value":"Literal"}, - {"value":"["}, - {"value":"2"}, - {"value":"]"}, - {"value":"]"} - ], - "position":{"character":21,"line":6}, - "textEdits":[{ - "newText":" -> tuple[Literal[1], Literal[2]]", - "range":{"end":{"character":21,"line":6},"start":{"character":21,"line":6}} - }] - }, - { - "label":[ - {"value":" -> "}, - {"value":"Literal"}, - {"value":"["}, - {"value":"0"}, - {"value":"]"} - ], - "position":{"character":15,"line":14}, - "textEdits":[{ - "newText":" -> Literal[0]", - "range":{"end":{"character":15,"line":14},"start":{"character":15,"line":14}} - }] - }])) + .expect_response_with(|result| { + let expected = [ + ExpectedInlayHint { + labels: &[ + " -> ", "tuple", "[", "Literal", "[", "1", "]", ", ", "Literal", "[", "2", + "]", "]", + ], + position: (6, 21), + text_edit: ExpectedTextEdit { + new_text: " -> tuple[Literal[1], Literal[2]]", + range_start: (6, 21), + range_end: (6, 21), + }, + }, + ExpectedInlayHint { + labels: &[" -> ", "Literal", "[", "0", "]"], + position: (14, 15), + text_edit: ExpectedTextEdit { + new_text: " -> Literal[0]", + range_start: (14, 15), + range_end: (14, 15), + }, + }, + ]; + inlay_hints_match_expected(result, &expected) + }) .unwrap(); interaction.shutdown().unwrap(); @@ -245,28 +213,21 @@ fn test_inlay_hint_disable_returns() { interaction .client .inlay_hint("inlay_hint_test.py", 0, 0, 100, 0) - .expect_response(json!([{ - "label":[ - {"value":": "}, - {"value":"tuple"}, - {"value":"["}, - {"value":"Literal"}, - {"value":"["}, - {"value":"1"}, - {"value":"]"}, - {"value":", "}, - {"value":"Literal"}, - {"value":"["}, - {"value":"2"}, - {"value":"]"}, - {"value":"]"} - ], - "position":{"character":6,"line":11}, - "textEdits":[{ - "newText":": tuple[Literal[1], Literal[2]]", - "range":{"end":{"character":6,"line":11},"start":{"character":6,"line":11}} - }] - }])) + .expect_response_with(|result| { + let expected = [ExpectedInlayHint { + labels: &[ + ": ", "tuple", "[", "Literal", "[", "1", "]", ", ", "Literal", "[", "2", "]", + "]", + ], + position: (11, 6), + text_edit: ExpectedTextEdit { + new_text: ": tuple[Literal[1], Literal[2]]", + range_start: (11, 6), + range_end: (11, 6), + }, + }]; + inlay_hints_match_expected(result, &expected) + }) .unwrap(); interaction.shutdown().unwrap(); @@ -325,3 +286,43 @@ fn test_inlay_hint_labels_support_goto_type_definition() { interaction.shutdown().unwrap(); } + +#[test] +fn test_inlay_hint_typing_literals_have_locations() { + let root = get_test_files_root(); + let mut interaction = LspInteraction::new(); + interaction.set_root(root.path().to_path_buf()); + interaction + .initialize(InitializeSettings { + configuration: Some(None), + ..Default::default() + }) + .unwrap(); + + interaction.client.did_open("inlay_hint_test.py"); + + interaction + .client + .inlay_hint("inlay_hint_test.py", 0, 0, 100, 0) + .expect_response_with(|result| { + let hints = match result { + Some(hints) => hints, + None => return false, + }; + + hints.iter().any(|hint| match &hint.label { + lsp_types::InlayHintLabel::LabelParts(parts) => parts.iter().any(|part| { + part.value == "Literal" + && part + .location + .as_ref() + .map(|loc| loc.uri.path().contains("typing.pyi")) + .unwrap_or(false) + }), + _ => false, + }) + }) + .unwrap(); + + interaction.shutdown().unwrap(); +} diff --git a/pyrefly/lib/test/lsp/lsp_interaction/notebook_inlay_hint.rs b/pyrefly/lib/test/lsp/lsp_interaction/notebook_inlay_hint.rs index 221fc198c..928f0c403 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/notebook_inlay_hint.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/notebook_inlay_hint.rs @@ -5,11 +5,12 @@ * LICENSE file in the root directory of this source tree. */ -use serde_json::json; - use crate::test::lsp::lsp_interaction::object_model::InitializeSettings; use crate::test::lsp::lsp_interaction::object_model::LspInteraction; +use crate::test::lsp::lsp_interaction::util::ExpectedInlayHint; +use crate::test::lsp::lsp_interaction::util::ExpectedTextEdit; use crate::test::lsp::lsp_interaction::util::get_test_files_root; +use crate::test::lsp::lsp_interaction::util::inlay_hints_match_expected; #[test] fn test_inlay_hints() { @@ -33,72 +34,56 @@ fn test_inlay_hints() { interaction .inlay_hint_cell("notebook.ipynb", "cell1", 0, 0, 100, 0) - .expect_response(json!([{ - "label": [ - {"value": " -> "}, - {"value": "tuple"}, - {"value": "["}, - {"value": "Literal"}, - {"value": "["}, - {"value": "1"}, - {"value": "]"}, - {"value": ", "}, - {"value": "Literal"}, - {"value": "["}, - {"value": "2"}, - {"value": "]"}, - {"value": "]"} - ], - "position": {"character": 21, "line": 0}, - "textEdits": [{ - "newText": " -> tuple[Literal[1], Literal[2]]", - "range": {"end": {"character": 21, "line": 0}, "start": {"character": 21, "line": 0}} - }] - }])) + .expect_response_with(|result| { + let expected = [ExpectedInlayHint { + labels: &[ + " -> ", "tuple", "[", "Literal", "[", "1", "]", ", ", "Literal", "[", "2", "]", + "]", + ], + position: (0, 21), + text_edit: ExpectedTextEdit { + new_text: " -> tuple[Literal[1], Literal[2]]", + range_start: (0, 21), + range_end: (0, 21), + }, + }]; + inlay_hints_match_expected(result, &expected) + }) .unwrap(); interaction .inlay_hint_cell("notebook.ipynb", "cell2", 0, 0, 100, 0) - .expect_response(json!([{ - "label": [ - {"value": ": "}, - {"value": "tuple"}, - {"value": "["}, - {"value": "Literal"}, - {"value": "["}, - {"value": "1"}, - {"value": "]"}, - {"value": ", "}, - {"value": "Literal"}, - {"value": "["}, - {"value": "2"}, - {"value": "]"}, - {"value": "]"} - ], - "position": {"character": 6, "line": 0}, - "textEdits": [{ - "newText": ": tuple[Literal[1], Literal[2]]", - "range": {"end": {"character": 6, "line": 0}, "start": {"character": 6, "line": 0}} - }] - }])) + .expect_response_with(|result| { + let expected = [ExpectedInlayHint { + labels: &[ + ": ", "tuple", "[", "Literal", "[", "1", "]", ", ", "Literal", "[", "2", "]", + "]", + ], + position: (0, 6), + text_edit: ExpectedTextEdit { + new_text: ": tuple[Literal[1], Literal[2]]", + range_start: (0, 6), + range_end: (0, 6), + }, + }]; + inlay_hints_match_expected(result, &expected) + }) .unwrap(); interaction .inlay_hint_cell("notebook.ipynb", "cell3", 0, 0, 100, 0) - .expect_response(json!([{ - "label": [ - {"value": " -> "}, - {"value": "Literal"}, - {"value": "["}, - {"value": "0"}, - {"value": "]"} - ], - "position": {"character": 15, "line": 0}, - "textEdits": [{ - "newText": " -> Literal[0]", - "range": {"end": {"character": 15, "line": 0}, "start": {"character": 15, "line": 0}} - }] - }])) + .expect_response_with(|result| { + let expected = [ExpectedInlayHint { + labels: &[" -> ", "Literal", "[", "0", "]"], + position: (0, 15), + text_edit: ExpectedTextEdit { + new_text: " -> Literal[0]", + range_start: (0, 15), + range_end: (0, 15), + }, + }]; + inlay_hints_match_expected(result, &expected) + }) .unwrap(); interaction.shutdown().unwrap(); } diff --git a/pyrefly/lib/test/lsp/lsp_interaction/util.rs b/pyrefly/lib/test/lsp/lsp_interaction/util.rs index 56af879ee..ace81cda2 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/util.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/util.rs @@ -8,6 +8,8 @@ use std::path::Path; use std::path::PathBuf; +use lsp_types::InlayHint; +use lsp_types::InlayHintLabel; use pyrefly_util::fs_anyhow; use tempfile::TempDir; @@ -35,6 +37,74 @@ pub fn bundled_typeshed_path() -> PathBuf { path } +pub struct ExpectedTextEdit<'a> { + pub new_text: &'a str, + pub range_start: (u32, u32), + pub range_end: (u32, u32), +} + +pub struct ExpectedInlayHint<'a> { + pub labels: &'a [&'a str], + pub position: (u32, u32), + pub text_edit: ExpectedTextEdit<'a>, +} + +pub fn inlay_hints_match_expected( + result: Option>, + expected: &[ExpectedInlayHint<'_>], +) -> bool { + let hints = match result { + Some(hints) => hints, + None => return false, + }; + if hints.len() != expected.len() { + return false; + } + + for (hint, expected_hint) in hints.iter().zip(expected.iter()) { + let position = hint.position; + if position.line != expected_hint.position.0 + || position.character != expected_hint.position.1 + { + return false; + } + + let parts = match &hint.label { + InlayHintLabel::LabelParts(parts) => parts, + _ => return false, + }; + if parts.len() != expected_hint.labels.len() { + return false; + } + for (part, expected_value) in parts.iter().zip(expected_hint.labels.iter()) { + if part.value != *expected_value { + return false; + } + if *expected_value == "Literal" && part.location.is_none() { + return false; + } + } + + let text_edits = match &hint.text_edits { + Some(edits) if edits.len() == 1 => &edits[0], + _ => return false, + }; + if text_edits.new_text != expected_hint.text_edit.new_text { + return false; + } + let start = &expected_hint.text_edit.range_start; + let end = &expected_hint.text_edit.range_end; + if text_edits.range.start.line != start.0 + || text_edits.range.start.character != start.1 + || text_edits.range.end.line != end.0 + || text_edits.range.end.character != end.1 + { + return false; + } + } + true +} + fn copy_dir_recursively(src: &Path, dst: &Path) { if !dst.exists() { std::fs::create_dir_all(dst).unwrap();