[ty] Enrich function argument auto-complete suggestions with annotated types

Author: RasmusNygren

crates/ty_ide/src/completion.rs

@@ -489,7 +489,7 @@ fn detect_function_arg_completions<'db>(
                 name,
                 qualified: None,
                 insert,
-                ty: None,
+                ty: p.ty,
                 kind: Some(CompletionKind::Variable),
                 module_name: None,
                 import: None,
@@ -3243,8 +3243,6 @@ abc(okay=1, ba<CURSOR> baz=5
         );
     }
 
-
-
     #[test]
     fn call_positional_only_argument() {
         // If the parameter is positional only we don't

crates/ty_ide/src/signature_help.rs

@@ -17,22 +17,24 @@ use ruff_text_size::{Ranged, TextRange, TextSize};
 use ty_python_semantic::ResolvedDefinition;
 use ty_python_semantic::SemanticModel;
 use ty_python_semantic::semantic_index::definition::Definition;
-use ty_python_semantic::types::ParameterKind;
 use ty_python_semantic::types::ide_support::{
     CallSignatureDetails, call_signature_details, find_active_signature_from_details,
 };
+use ty_python_semantic::types::{ParameterKind, Type};
 
 // TODO: We may want to add special-case handling for calls to constructors
 // so the class docstring is used in place of (or inaddition to) any docstring
 // associated with the __new__ or __init__ call.
 
 /// Information about a function parameter
 #[derive(Debug, Clone, PartialEq, Eq)]
-pub struct ParameterDetails {
+pub struct ParameterDetails<'db> {
     /// The parameter name (e.g., "param1")
     pub name: String,
     /// The parameter label in the signature (e.g., "param1: str")
     pub label: String,
+    /// The annotated type of the parameter, if any
+    pub ty: Option<Type<'db>>,
     /// Documentation specific to the parameter, typically extracted from the
     /// function's docstring
     pub documentation: Option<String>,
@@ -42,32 +44,32 @@ pub struct ParameterDetails {
 
 /// Information about a function signature
 #[derive(Debug, Clone, PartialEq, Eq)]
-pub struct SignatureDetails {
+pub struct SignatureDetails<'db> {
     /// Text representation of the full signature (including input parameters and return type).
     pub label: String,
     /// Documentation for the signature, typically from the function's docstring.
     pub documentation: Option<Docstring>,
     /// Information about each of the parameters in left-to-right order.
-    pub parameters: Vec<ParameterDetails>,
+    pub parameters: Vec<ParameterDetails<'db>>,
     /// Index of the parameter that corresponds to the argument where the
     /// user's cursor is currently positioned.
     pub active_parameter: Option<usize>,
 }
 
 /// Signature help information for function calls
 #[derive(Debug, Clone, PartialEq, Eq)]
-pub struct SignatureHelpInfo {
+pub struct SignatureHelpInfo<'db> {
     /// Information about each of the signatures for the function call. We
     /// need to handle multiple because of unions, overloads, and composite
     /// calls like constructors (which invoke both __new__ and __init__).
-    pub signatures: Vec<SignatureDetails>,
+    pub signatures: Vec<SignatureDetails<'db>>,
     /// Index of the "active signature" which is the first signature where
     /// all arguments that are currently present in the code map to parameters.
     pub active_signature: Option<usize>,
 }
 
 /// Signature help information for function calls at the given position
-pub fn signature_help(db: &dyn Db, file: File, offset: TextSize) -> Option<SignatureHelpInfo> {
+pub fn signature_help(db: &dyn Db, file: File, offset: TextSize) -> Option<SignatureHelpInfo<'_>> {
     let parsed = parsed_module(db, file).load(db);
 
     // Get the call expression at the given position.
@@ -169,11 +171,11 @@ fn get_argument_index(call_expr: &ast::ExprCall, offset: TextSize) -> usize {
 }
 
 /// Create signature details from `CallSignatureDetails`.
-fn create_signature_details_from_call_signature_details(
+fn create_signature_details_from_call_signature_details<'db>(
     db: &dyn crate::Db,
-    details: &CallSignatureDetails,
+    details: &CallSignatureDetails<'db>,
     current_arg_index: usize,
-) -> SignatureDetails {
+) -> SignatureDetails<'db> {
     let signature_label = details.label.clone();
 
     let documentation = get_callable_documentation(db, details.definition);
@@ -204,6 +206,7 @@ fn create_signature_details_from_call_signature_details(
         documentation.as_ref(),
         &details.parameter_names,
         &details.parameter_kinds,
+        &details.parameter_types,
     );
     SignatureDetails {
         label: signature_label,
@@ -222,13 +225,14 @@ fn get_callable_documentation(
 }
 
 /// Create `ParameterDetails` objects from parameter label offsets.
-fn create_parameters_from_offsets(
+fn create_parameters_from_offsets<'db>(
     parameter_offsets: &[TextRange],
     signature_label: &str,
     docstring: Option<&Docstring>,
     parameter_names: &[String],
     parameter_kinds: &[ParameterKind],
-) -> Vec<ParameterDetails> {
+    parameter_types: &[Option<Type<'db>>],
+) -> Vec<ParameterDetails<'db>> {
     // Extract parameter documentation from the function's docstring if available.
     let param_docs = if let Some(docstring) = docstring {
         docstring.parameter_documentation()
@@ -254,10 +258,12 @@ fn create_parameters_from_offsets(
                 parameter_kinds.get(i),
                 Some(ParameterKind::PositionalOnly { .. })
             );
+            let ty = parameter_types.get(i).copied().flatten();
 
             ParameterDetails {
                 name: param_name.to_string(),
                 label,
+                ty,
                 documentation: param_docs.get(param_name).cloned(),
                 is_positional_only,
             }
@@ -1183,7 +1189,7 @@ def ab(a: int, *, c: int):
     }
 
     impl CursorTest {
-        fn signature_help(&self) -> Option<SignatureHelpInfo> {
+        fn signature_help(&self) -> Option<SignatureHelpInfo<'_>> {
             crate::signature_help::signature_help(&self.db, self.cursor.file, self.cursor.offset)
         }
 

crates/ty_python_semantic/src/types/ide_support.rs

@@ -462,6 +462,9 @@ pub struct CallSignatureDetails<'db> {
     /// Parameter kinds, useful to determine correct autocomplete suggestions.
     pub parameter_kinds: Vec<ParameterKind<'db>>,
 
+    /// Parameter kinds, useful to determine correct autocomplete suggestions.
+    pub parameter_types: Vec<Option<Type<'db>>>,
+
     /// The definition where this callable was originally defined (useful for
     /// extracting docstrings).
     pub definition: Option<Definition<'db>>,
@@ -520,11 +523,12 @@ pub fn call_signature_details<'db>(
                 let display_details = signature.display(model.db()).to_string_parts();
                 let parameter_label_offsets = display_details.parameter_ranges;
                 let parameter_names = display_details.parameter_names;
-                let parameter_kinds = signature
-                    .parameters()
-                    .iter()
-                    .map(|param| param.kind().clone())
-                    .collect();
+                let (parameter_kinds, parameter_types): (Vec<ParameterKind>, Vec<Option<Type>>) =
+                    signature
+                        .parameters()
+                        .iter()
+                        .map(|param| (param.kind().clone(), param.annotated_type()))
+                        .unzip();
 
                 CallSignatureDetails {
                     definition: signature.definition(),
@@ -533,6 +537,7 @@ pub fn call_signature_details<'db>(
                     parameter_label_offsets,
                     parameter_names,
                     parameter_kinds,
+                    parameter_types,
                     argument_to_parameter_mapping,
                 }
             })