Track positional and keyword arguments

Author: aaron-ang

pyrefly/lib/state/lsp.rs

@@ -86,6 +86,13 @@ use crate::types::callable::Params;
 use crate::types::module::ModuleType;
 use crate::types::types::Type;
 
+#[derive(Debug, Clone)]
+enum ActiveArgument {
+    Positional(usize),
+    Keyword(Name),
+    Next(usize),
+}
+
 fn default_true() -> bool {
     true
 }
@@ -844,25 +851,47 @@ impl<'a> Transaction<'a> {
     fn visit_finding_signature_range(
         x: &Expr,
         find: TextSize,
-        res: &mut Option<(TextRange, TextRange, usize)>,
+        res: &mut Option<(TextRange, TextRange, ActiveArgument)>,
     ) {
         if let Expr::Call(call) = x
             && call.arguments.range.contains_inclusive(find)
         {
+            // Check positional arguments
             for (i, arg) in call.arguments.args.as_ref().iter().enumerate() {
                 if arg.range().contains_inclusive(find) {
                     Self::visit_finding_signature_range(arg, find, res);
                     if res.is_some() {
                         return;
                     }
-                    *res = Some((call.func.range(), call.arguments.range, i));
+                    *res = Some((
+                        call.func.range(),
+                        call.arguments.range,
+                        ActiveArgument::Positional(i),
+                    ));
+                    return;
+                }
+            }
+            // Check keyword arguments
+            let positional_count = call.arguments.args.len();
+            for (j, kw) in call.arguments.keywords.iter().enumerate() {
+                if kw.range.contains_inclusive(find) {
+                    Self::visit_finding_signature_range(&kw.value, find, res);
+                    if res.is_some() {
+                        return;
+                    }
+                    let active_argument = match kw.arg.as_ref() {
+                        Some(identifier) => ActiveArgument::Keyword(identifier.id.clone()),
+                        None => ActiveArgument::Positional(positional_count + j),
+                    };
+                    *res = Some((call.func.range(), call.arguments.range, active_argument));
+                    return;
                 }
             }
             if res.is_none() {
                 *res = Some((
                     call.func.range(),
                     call.arguments.range,
-                    call.arguments.len(),
+                    ActiveArgument::Next(call.arguments.len()),
                 ));
             }
         } else {
@@ -875,11 +904,17 @@ impl<'a> Transaction<'a> {
         &self,
         handle: &Handle,
         position: TextSize,
-    ) -> Option<(Vec<Type>, usize, usize)> {
+    ) -> Option<(Vec<Type>, usize, ActiveArgument)> {
         let mod_module = self.get_ast(handle)?;
         let mut res = None;
         mod_module.visit(&mut |x| Self::visit_finding_signature_range(x, position, &mut res));
-        let (callee_range, call_args_range, arg_index) = res?;
+        let (callee_range, call_args_range, mut active_argument) = res?;
+        if let ActiveArgument::Next(index) = &mut active_argument
+            && let Some(next_index) =
+                self.count_argument_separators_before(handle, call_args_range, position)
+        {
+            *index = next_index;
+        }
         let answers = self.get_answers(handle)?;
         if let Some((overloads, chosen_overload_index)) =
             answers.get_all_overload_trace(call_args_range)
@@ -888,12 +923,12 @@ impl<'a> Transaction<'a> {
             Some((
                 callables,
                 chosen_overload_index.unwrap_or_default(),
-                arg_index,
+                active_argument,
             ))
         } else {
             answers
                 .get_type_trace(callee_range)
-                .map(|t| (vec![t], 0, arg_index))
+                .map(|t| (vec![t], 0, active_argument))
         }
     }
 
@@ -903,27 +938,33 @@ impl<'a> Transaction<'a> {
         position: TextSize,
     ) -> Option<SignatureHelp> {
         self.get_callables_from_call(handle, position).map(
-            |(callables, chosen_overload_index, arg_index)| SignatureHelp {
-                signatures: callables
+            |(callables, chosen_overload_index, active_argument)| {
+                let signatures = callables
                     .into_iter()
-                    .map(|t| Self::create_signature_information(t, arg_index))
-                    .collect_vec(),
-                active_signature: Some(chosen_overload_index as u32),
-                active_parameter: Some(arg_index as u32),
+                    .map(|t| Self::create_signature_information(t, &active_argument))
+                    .collect_vec();
+                let active_parameter = signatures
+                    .get(chosen_overload_index)
+                    .and_then(|info| info.active_parameter);
+                SignatureHelp {
+                    signatures,
+                    active_signature: Some(chosen_overload_index as u32),
+                    active_parameter,
+                }
             },
         )
     }
 
-    fn create_signature_information(type_: Type, arg_index: usize) -> SignatureInformation {
+    fn create_signature_information(
+        type_: Type,
+        active_argument: &ActiveArgument,
+    ) -> SignatureInformation {
         let type_ = type_.deterministic_printing();
         let label = type_.as_hover_string();
         let (parameters, active_parameter) =
             if let Some(params) = Self::normalize_singleton_function_type_into_params(type_) {
-                let active_parameter = if arg_index < params.len() {
-                    Some(arg_index as u32)
-                } else {
-                    None
-                };
+                let active_parameter =
+                    Self::active_parameter_index(&params, active_argument).map(|idx| idx as u32);
                 (
                     Some(params.map(|param| ParameterInformation {
                         label: ParameterLabel::Simple(format!("{param}")),
@@ -942,6 +983,46 @@ impl<'a> Transaction<'a> {
         }
     }
 
+    fn active_parameter_index(params: &[Param], active_argument: &ActiveArgument) -> Option<usize> {
+        match active_argument {
+            ActiveArgument::Positional(index) | ActiveArgument::Next(index) => {
+                (*index < params.len()).then_some(*index)
+            }
+            ActiveArgument::Keyword(name) => params.iter().position(|param| {
+                Self::parameter_name(param).is_some_and(|param_name| param_name == name)
+            }),
+        }
+    }
+
+    fn parameter_name(param: &Param) -> Option<&Name> {
+        match param {
+            Param::PosOnly(Some(name), ..)
+            | Param::Pos(name, ..)
+            | Param::VarArg(Some(name), ..)
+            | Param::KwOnly(name, ..)
+            | Param::Kwargs(Some(name), ..) => Some(name),
+            _ => None,
+        }
+    }
+
+    fn count_argument_separators_before(
+        &self,
+        handle: &Handle,
+        arguments_range: TextRange,
+        position: TextSize,
+    ) -> Option<usize> {
+        let module = self.get_module_info(handle)?;
+        let contents = module.contents();
+        let start = arguments_range.start().to_usize();
+        let end = arguments_range.end().to_usize().min(contents.len());
+        if start >= end {
+            return Some(0);
+        }
+        let pos = position.to_usize().clamp(start, end);
+        let slice = &contents[start..pos];
+        Some(slice.bytes().filter(|&b| b == b',').count())
+    }
+
     fn normalize_singleton_function_type_into_params(type_: Type) -> Option<Vec<Param>> {
         let callable = type_.to_callable()?;
         // We will drop the self parameter for signature help
@@ -2169,11 +2250,12 @@ impl<'a> Transaction<'a> {
         position: TextSize,
         completions: &mut Vec<CompletionItem>,
     ) {
-        if let Some((callables, chosen_overload_index, arg_index)) =
+        if let Some((callables, chosen_overload_index, active_argument)) =
             self.get_callables_from_call(handle, position)
             && let Some(callable) = callables.get(chosen_overload_index)
             && let Some(params) =
                 Self::normalize_singleton_function_type_into_params(callable.clone())
+            && let Some(arg_index) = Self::active_parameter_index(&params, &active_argument)
             && let Some(param) = params.get(arg_index)
         {
             Self::add_literal_completions_from_type(param.as_type(), completions);

pyrefly/lib/test/lsp/signature_help.rs

@@ -128,6 +128,99 @@ Signature Help Result: active=0
     );
 }
 
+#[test]
+fn positional_arguments_test() {
+    let code = r#"
+def f(x: int, y: int, z: int) -> None: ...
+
+f(1,,)
+#   ^
+f(1,,)
+#    ^
+f(1,,3)
+#   ^
+"#;
+    let report = get_batched_lsp_operations_report_allow_error(&[("main", code)], get_test_report);
+    assert_eq!(
+        r#"
+# main.py
+4 | f(1,,)
+        ^
+Signature Help Result: active=0
+- def f(
+    x: int,
+    y: int,
+    z: int
+) -> None, parameters=[x: int, y: int, z: int], active parameter = 1
+
+6 | f(1,,)
+         ^
+Signature Help Result: active=0
+- def f(
+    x: int,
+    y: int,
+    z: int
+) -> None, parameters=[x: int, y: int, z: int], active parameter = 2
+
+8 | f(1,,3)
+        ^
+Signature Help Result: active=0
+- def f(
+    x: int,
+    y: int,
+    z: int
+) -> None, parameters=[x: int, y: int, z: int], active parameter = 1
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+
+#[test]
+fn keyword_arguments_test() {
+    let code = r#"
+def f(a: str, b: int) -> None: ...
+
+f(a)
+# ^
+f(a=)
+#  ^
+f(b=)
+#  ^
+"#;
+    let report = get_batched_lsp_operations_report_allow_error(&[("main", code)], get_test_report);
+    assert_eq!(
+        r#"
+# main.py
+4 | f(a)
+      ^
+Signature Help Result: active=0
+- def f(
+    a: str,
+    b: int
+) -> None, parameters=[a: str, b: int], active parameter = 0
+
+6 | f(a=)
+       ^
+Signature Help Result: active=0
+- def f(
+    a: str,
+    b: int
+) -> None, parameters=[a: str, b: int], active parameter = 0
+
+8 | f(b=)
+       ^
+Signature Help Result: active=0
+- def f(
+    a: str,
+    b: int
+) -> None, parameters=[a: str, b: int], active parameter = 1
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+
 #[test]
 fn simple_incomplete_function_call_test() {
     let code = r#"