[ty] Add autocomplete suggestions for function arguments

This adds autocomplete suggestions for function arguments. For example,
`okay` in:

```python
def foo(okay=None):

foo(o<CURSOR>
```

This also ensures that we don't suggest a keyword argument if it has
already been used.

Closes astral-sh/issues#1550
This commit is contained in:
Rasmus Nygren 2025-12-03 21:58:01 +01:00 committed by Andrew Gallant
parent 2d3466eccf
commit eac8a90cc4
5 changed files with 268 additions and 10 deletions

View File

@ -9,6 +9,7 @@ use ruff_python_ast::token::{Token, TokenAt, TokenKind, Tokens};
use ruff_python_ast::{self as ast, AnyNodeRef};
use ruff_python_codegen::Stylist;
use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
use rustc_hash::FxHashSet;
use ty_python_semantic::types::UnionType;
use ty_python_semantic::{
Completion as SemanticCompletion, KnownModule, ModuleName, NameKind, SemanticModel,
@ -20,7 +21,7 @@ use crate::find_node::covering_node;
use crate::goto::Definitions;
use crate::importer::{ImportRequest, Importer};
use crate::symbols::QueryPattern;
use crate::{Db, all_symbols};
use crate::{Db, all_symbols, signature_help};
/// A collection of completions built up from various sources.
#[derive(Clone)]
@ -436,6 +437,10 @@ pub fn completion<'db>(
);
}
}
if let Some(arg_completions) = detect_function_arg_completions(db, file, &parsed, offset) {
completions.extend(arg_completions);
}
}
if is_raising_exception(tokens) {
@ -451,10 +456,89 @@ pub fn completion<'db>(
!ty.is_notimplemented(db)
});
}
completions.into_completions()
}
/// Detect and construct completions for unset function arguments.
///
/// Suggestions are only provided if the cursor is currently inside a
/// function call and the function arguments have not 1) already been
/// set and 2) been defined as positional-only.
fn detect_function_arg_completions<'db>(
db: &'db dyn Db,
file: File,
parsed: &ParsedModuleRef,
offset: TextSize,
) -> Option<Vec<Completion<'db>>> {
let sig_help = signature_help(db, file, offset)?;
let set_function_args = detect_set_function_args(parsed, offset);
let completions = sig_help
.signatures
.iter()
.flat_map(|sig| &sig.parameters)
.filter(|p| !p.is_positional_only && !set_function_args.contains(&p.name.as_str()))
.map(|p| {
let name = Name::new(&p.name);
let documentation = p
.documentation
.as_ref()
.map(|d| Docstring::new(d.to_owned()));
let insert = Some(format!("{name}=").into_boxed_str());
Completion {
name,
qualified: None,
insert,
ty: None,
kind: Some(CompletionKind::Variable),
module_name: None,
import: None,
builtin: false,
is_type_check_only: false,
is_definitively_raisable: false,
documentation,
}
})
.collect();
Some(completions)
}
/// Returns function arguments that have already been set.
///
/// If `offset` is inside an arguments node, this returns
/// the list of argument names that are already set.
///
/// For example, given:
///
/// ```python
/// def abc(foo, bar, baz): ...
/// abc(foo=1, bar=2, b<CURSOR>)
/// ```
///
/// the resulting value is `["foo", "bar"]`
///
/// This is useful to be able to exclude autocomplete suggestions
/// for arguments that have already been set to some value.
///
/// If the parent node is not an arguments node, the return value
/// is an empty Vec.
fn detect_set_function_args(parsed: &ParsedModuleRef, offset: TextSize) -> FxHashSet<&str> {
let range = TextRange::empty(offset);
covering_node(parsed.syntax().into(), range)
.parent()
.and_then(|node| match node {
ast::AnyNodeRef::Arguments(args) => Some(args),
_ => None,
})
.map(|args| {
args.keywords
.iter()
.filter_map(|kw| kw.arg.as_ref().map(|ident| ident.id.as_str()))
.collect()
})
.unwrap_or_default()
}
pub(crate) struct ImportEdit {
pub label: String,
pub edit: Edit,
@ -2386,10 +2470,11 @@ def frob(): ...
",
);
// FIXME: Should include `foo`.
assert_snapshot!(
builder.skip_keywords().skip_builtins().build().snapshot(),
@"<No completions found after filtering out completions>",
@r"
foo
",
);
}
@ -2401,10 +2486,11 @@ def frob(): ...
",
);
// FIXME: Should include `foo`.
assert_snapshot!(
builder.skip_keywords().skip_builtins().build().snapshot(),
@"<No completions found after filtering out completions>",
@r"
foo
",
);
}
@ -3039,7 +3125,6 @@ quux.<CURSOR>
");
}
// We don't yet take function parameters into account.
#[test]
fn call_prefix1() {
let builder = completion_test_builder(
@ -3052,7 +3137,159 @@ bar(o<CURSOR>
",
);
assert_snapshot!(builder.skip_keywords().skip_builtins().build().snapshot(), @"foo");
assert_snapshot!(
builder.skip_keywords().skip_builtins().build().snapshot(),
@r"
foo
okay
"
);
}
#[test]
fn call_keyword_only_argument() {
let builder = completion_test_builder(
"\
def bar(*, okay): ...
foo = 1
bar(o<CURSOR>
",
);
assert_snapshot!(
builder.skip_keywords().skip_builtins().build().snapshot(),
@r"
foo
okay
"
);
}
#[test]
fn call_multiple_keyword_arguments() {
let builder = completion_test_builder(
"\
def foo(bar, baz, barbaz): ...
foo(b<CURSOR>
",
);
assert_snapshot!(
builder.skip_keywords().skip_builtins().build().snapshot(),
@r"
bar
barbaz
baz
"
);
}
#[test]
fn call_multiple_keyword_arguments_some_set() {
let builder = completion_test_builder(
"\
def foo(bar, baz): ...
foo(bar=1, b<CURSOR>
",
);
assert_snapshot!(
builder.skip_keywords().skip_builtins().build().snapshot(),
@r"
baz
"
);
}
#[test]
fn call_arguments_multi_def() {
let builder = completion_test_builder(
"\
def abc(okay, x): ...
def bar(not_okay, y): ...
def baz(foobarbaz, z): ...
abc(o<CURSOR>
",
);
assert_snapshot!(
builder.skip_keywords().skip_builtins().build().snapshot(),
@r"
okay
"
);
}
#[test]
fn call_arguments_cursor_middle() {
let builder = completion_test_builder(
"\
def abc(okay, foo, bar, baz): ...
abc(okay=1, ba<CURSOR> baz=5
",
);
assert_snapshot!(
builder.skip_keywords().skip_builtins().build().snapshot(),
@r"
bar
"
);
}
#[test]
fn call_positional_only_argument() {
// If the parameter is positional only we don't
// want to suggest it as specifying by name
// is not valid.
let builder = completion_test_builder(
"\
def bar(okay, /): ...
foo = 1
bar(o<CURSOR>
",
);
assert_snapshot!(
builder.skip_keywords().skip_builtins().build().snapshot(),
@"foo"
);
}
#[test]
fn call_positional_only_keyword_only_argument_mix() {
// If the parameter is positional only we don't
// want to suggest it as specifying by name
// is not valid.
let builder = completion_test_builder(
"\
def bar(not_okay, no, /, okay, *, okay_abc, okay_okay): ...
foo = 1
bar(o<CURSOR>
",
);
assert_snapshot!(
builder.skip_keywords().skip_builtins().build().snapshot(),
@r"
foo
okay
okay_abc
okay_okay
"
);
}
#[test]
@ -3070,6 +3307,7 @@ bar(<CURSOR>
assert_snapshot!(builder.skip_keywords().skip_builtins().build().snapshot(), @r"
bar
foo
okay
");
}

View File

@ -17,6 +17,7 @@ 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,
};
@ -35,6 +36,8 @@ pub struct ParameterDetails {
/// Documentation specific to the parameter, typically extracted from the
/// function's docstring
pub documentation: Option<String>,
/// True if the parameter is positional-only.
pub is_positional_only: bool,
}
/// Information about a function signature
@ -200,6 +203,7 @@ fn create_signature_details_from_call_signature_details(
&signature_label,
documentation.as_ref(),
&details.parameter_names,
&details.parameter_kinds,
);
SignatureDetails {
label: signature_label,
@ -223,6 +227,7 @@ fn create_parameters_from_offsets(
signature_label: &str,
docstring: Option<&Docstring>,
parameter_names: &[String],
parameter_kinds: &[ParameterKind],
) -> Vec<ParameterDetails> {
// Extract parameter documentation from the function's docstring if available.
let param_docs = if let Some(docstring) = docstring {
@ -245,11 +250,16 @@ fn create_parameters_from_offsets(
// Get the parameter name for documentation lookup.
let param_name = parameter_names.get(i).map(String::as_str).unwrap_or("");
let is_positional_only = matches!(
parameter_kinds.get(i),
Some(ParameterKind::PositionalOnly { .. })
);
ParameterDetails {
name: param_name.to_string(),
label,
documentation: param_docs.get(param_name).cloned(),
is_positional_only,
}
})
.collect()

View File

@ -30,6 +30,7 @@ pub(crate) use self::infer::{
TypeContext, infer_deferred_types, infer_definition_types, infer_expression_type,
infer_expression_types, infer_scope_types, static_expression_truthiness,
};
pub use self::signatures::ParameterKind;
pub(crate) use self::signatures::{CallableSignature, Signature};
pub(crate) use self::subclass_of::{SubclassOfInner, SubclassOfType};
pub use crate::diagnostic::add_inferred_python_version_hint_to_diagnostic;

View File

@ -6,7 +6,7 @@ use crate::semantic_index::definition::Definition;
use crate::semantic_index::definition::DefinitionKind;
use crate::semantic_index::{attribute_scopes, global_scope, semantic_index, use_def_map};
use crate::types::call::{CallArguments, MatchedArgument};
use crate::types::signatures::Signature;
use crate::types::signatures::{ParameterKind, Signature};
use crate::types::{CallDunderError, UnionType};
use crate::types::{CallableTypes, ClassBase, KnownClass, Type, TypeContext};
use crate::{Db, DisplaySettings, HasType, SemanticModel};
@ -459,6 +459,9 @@ pub struct CallSignatureDetails<'db> {
/// This provides easy access to parameter names for documentation lookup.
pub parameter_names: Vec<String>,
/// Parameter kinds, useful to determine correct autocomplete suggestions.
pub parameter_kinds: Vec<ParameterKind<'db>>,
/// The definition where this callable was originally defined (useful for
/// extracting docstrings).
pub definition: Option<Definition<'db>>,
@ -517,6 +520,11 @@ 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();
CallSignatureDetails {
definition: signature.definition(),
@ -524,6 +532,7 @@ pub fn call_signature_details<'db>(
label: display_details.label,
parameter_label_offsets,
parameter_names,
parameter_kinds,
argument_to_parameter_mapping,
}
})

View File

@ -2292,7 +2292,7 @@ impl<'db> Parameter<'db> {
}
#[derive(Clone, Debug, PartialEq, Eq, Hash, salsa::Update, get_size2::GetSize)]
pub(crate) enum ParameterKind<'db> {
pub enum ParameterKind<'db> {
/// Positional-only parameter, e.g. `def f(x, /): ...`
PositionalOnly {
/// Parameter name.