From 0edd97dd41bd0529f439a38d41831fa68866d201 Mon Sep 17 00:00:00 2001 From: RasmusNygren Date: Tue, 30 Dec 2025 14:10:56 +0100 Subject: [PATCH] [ty] Add autocomplete suggestions for class arguments (#22110) --- crates/ty_ide/src/completion.rs | 315 +++++++++++++++++-- crates/ty_python_semantic/src/types.rs | 2 +- crates/ty_python_semantic/src/types/class.rs | 2 +- 3 files changed, 283 insertions(+), 36 deletions(-) diff --git a/crates/ty_ide/src/completion.rs b/crates/ty_ide/src/completion.rs index 7f08e1390f..593b1008bf 100644 --- a/crates/ty_ide/src/completion.rs +++ b/crates/ty_ide/src/completion.rs @@ -12,6 +12,7 @@ use ruff_python_codegen::Stylist; use ruff_text_size::{Ranged, TextRange, TextSize}; use rustc_hash::FxHashSet; use ty_module_resolver::{KnownModule, ModuleName}; +use ty_python_semantic::HasType; use ty_python_semantic::types::UnionType; use ty_python_semantic::{ Completion as SemanticCompletion, NameKind, SemanticModel, @@ -60,6 +61,7 @@ pub fn completion<'db>( completions.extend(semantic_completions); if scoped.is_some() { add_keyword_completions(db, &mut completions); + add_argument_completions(db, &model, &context.cursor, &mut completions); } if settings.auto_import { if let Some(scoped) = scoped { @@ -75,8 +77,6 @@ pub fn completion<'db>( ); } } - - add_function_arg_completions(db, file, &context.cursor, &mut completions); } } @@ -419,6 +419,25 @@ impl<'db> Completion<'db> { } } + fn argument(name: &str, ty: Option>, documentation: Option<&str>) -> Self { + let insert = Some(format!("{name}=").into_boxed_str()); + let documentation = documentation.map(|d| Docstring::new(d.to_owned())); + + Completion { + name: name.into(), + qualified: None, + insert, + ty, + kind: Some(CompletionKind::Variable), + module_name: None, + import: None, + builtin: false, + is_type_check_only: false, + is_definitively_raisable: false, + documentation, + } + } + /// Returns true when this completion refers to the /// `NotImplemented` builtin. fn is_notimplemented(&self, db: &dyn Db) -> bool { @@ -1063,7 +1082,77 @@ enum Sort { Lower, } -/// Detect and construct completions for unset function arguments. +/// Detect and add completions for unset arguments. +fn add_argument_completions<'db>( + db: &'db dyn Db, + model: &SemanticModel<'db>, + cursor: &ContextCursor<'_>, + completions: &mut Completions<'db>, +) { + for node in cursor.covering_node(cursor.range).ancestors() { + match node { + ast::AnyNodeRef::ExprCall(call) => { + if call.arguments.range().contains_range(cursor.range) { + add_function_arg_completions(db, model.file(), cursor, completions); + } + return; + } + ast::AnyNodeRef::StmtClassDef(class_def) => { + if let Some(arguments) = class_def.arguments.as_deref() + && arguments.range().contains_range(cursor.range) + { + add_class_arg_completions(model, class_def, completions); + } + return; + } + node => { + if node.is_statement() { + return; + } + } + } + } +} + +/// Detect and add completions for unset class arguments. +/// +/// Some arguments we know are always valid and thus they are easy +/// to provide. The `metaclass` keyword is always valid. +/// For `typing.TypedDict` subclasses, we add +/// `TypedDict` specific keywords like `total`. +fn add_class_arg_completions<'db>( + model: &SemanticModel<'db>, + class_def: &ast::StmtClassDef, + completions: &mut Completions<'db>, +) { + let is_set = |name| { + class_def + .arguments + .as_ref() + .is_some_and(|args| args.find_keyword(name).is_some()) + }; + + if !is_set("metaclass") { + let ty = Some(KnownClass::Type.to_subclass_of(model.db())); + completions.add(Completion::argument("metaclass", ty, None)); + } + + let is_typed_dict = class_def + .inferred_type(model) + .and_then(Type::as_class_literal) + .is_some_and(|t| t.is_typed_dict(model.db())); + + // TODO: Handle PEP 728 that adds two extra keywords, + // closed and extra_items. + // + // See https://peps.python.org/pep-0728/ + if is_typed_dict && !is_set("total") { + let ty = Some(KnownClass::Bool.to_instance(model.db())); + completions.add(Completion::argument("total", ty, None)); + } +} + +/// Detect and add 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 @@ -1074,18 +1163,15 @@ fn add_function_arg_completions<'db>( cursor: &ContextCursor<'_>, completions: &mut Completions<'db>, ) { - // But be careful: this isn't as simple as just finding a call - // expression. We also have to make sure we are in the "arguments" - // portion of the call. Otherwise we risk incorrectly returning - // something for `()(arg1, arg2)`-style expressions. - if !cursor - .covering_node(TextRange::empty(cursor.offset)) - .ancestors() - .take_while(|node| !node.is_statement()) - .any(|node| node.is_arguments()) - { - return; - } + debug_assert!( + cursor + .covering_node(cursor.range) + .ancestors() + .take_while(|node| !node.is_statement()) + .any(|node| node.is_arguments()), + "Should only be called if we're already certain we're in an arguments node to avoid \ + adding completions for something like `()(arg1, arg2)`-style expressions" + ); let Some(sig_help) = signature_help(db, file, cursor.offset) else { return; @@ -1098,25 +1184,11 @@ fn add_function_arg_completions<'db>( continue; } - 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()); - completions.add(Completion { - name, - qualified: None, - insert, - ty: p.ty, - kind: Some(CompletionKind::Variable), - module_name: None, - import: None, - builtin: false, - is_type_check_only: false, - is_definitively_raisable: false, - documentation, - }); + completions.add(Completion::argument( + &p.name, + p.ty, + p.documentation.as_deref(), + )); } } } @@ -3032,6 +3104,7 @@ class Foo(): assert_snapshot!(builder.skip_keywords().skip_builtins().build().snapshot(), @r" Bar Foo + metaclass= "); } @@ -3049,6 +3122,7 @@ class Bar: ... assert_snapshot!(builder.skip_keywords().skip_builtins().build().snapshot(), @r" Bar Foo + metaclass= "); } @@ -3066,6 +3140,7 @@ class Bar: ... assert_snapshot!(builder.skip_keywords().skip_builtins().build().snapshot(), @r" Bar Foo + metaclass= "); } @@ -3081,9 +3156,159 @@ class Foo(", assert_snapshot!(builder.skip_keywords().skip_builtins().build().snapshot(), @r" Bar Foo + metaclass= "); } + #[test] + fn class_metaclass() { + let builder = completion_test_builder( + "\ +class Foo(meta", + ); + + builder + .skip_keywords() + .skip_builtins() + .build() + .contains("metaclass"); + } + + #[test] + fn class_metaclass_set() { + let builder = completion_test_builder( + "\ +class Foo(metaclass=x, meta", + ); + + builder + .skip_keywords() + .skip_builtins() + .build() + .not_contains("metaclass"); + } + + #[test] + fn class_metaclass_generic() { + let builder = completion_test_builder( + "\ +class Foo[T](meta", + ); + + builder + .skip_keywords() + .skip_builtins() + .build() + .contains("metaclass"); + } + + #[test] + fn class_typed_dict_total() { + let builder = completion_test_builder( + "\ +from typing import TypedDict + +class Foo(TypedDict, tot +", + ); + + builder + .skip_keywords() + .skip_builtins() + .build() + .contains("total"); + } + + #[test] + fn class_typed_dict_total_alias() { + let builder = completion_test_builder( + "\ +from typing import TypedDict as TD + +class Foo(TD, tot +", + ); + + builder + .skip_keywords() + .skip_builtins() + .build() + .contains("total"); + } + + #[test] + fn class_typed_dict_total_set() { + let builder = completion_test_builder( + "\ +from typing import TypedDict + +class Foo(TypedDict, total=False, tot +", + ); + + builder + .skip_keywords() + .skip_builtins() + .build() + .not_contains("total"); + } + + #[test] + fn class_typed_dict_total_subclass() { + let builder = completion_test_builder( + "\ +from typing import TypedDict + +class Foo(TypedDict): + x: int + +class Bar(Foo, to) +", + ); + + builder + .skip_keywords() + .skip_builtins() + .build() + .contains("total"); + } + + #[test] + fn class_typed_dict_total_pep695_generic() { + let builder = completion_test_builder( + "\ +from typing import TypedDict + +class Foo[T](TypedDict, to) +", + ); + + builder + .skip_keywords() + .skip_builtins() + .build() + .contains("total"); + } + + #[test] + fn class_typed_dict_total_typevar_generic() { + let builder = completion_test_builder( + "\ +from typing import Generic, TypeVar, TypedDict + +T = TypeVar('T') + +class Foo(TypedDict, Generic[T], to) +", + ); + + builder + .skip_keywords() + .skip_builtins() + .build() + .contains("total"); + } + #[test] fn class_init1() { let builder = completion_test_builder( @@ -3756,6 +3981,28 @@ bar( "); } + #[test] + fn call_attribute_argument_no_arg_completions() { + let builder = completion_test_builder( + "\ +class A: + class B: + class C: ... + +def f(aaaa): ... + +f(A.B.) +", + ); + + builder + .skip_keywords() + .skip_builtins() + .build() + .contains("C") + .not_contains("aaaa"); + } + #[test] fn duplicate1() { let builder = completion_test_builder( diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index 13d62425b5..b76d6ab520 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -1257,7 +1257,7 @@ impl<'db> Type<'db> { } } - pub(crate) const fn as_class_literal(self) -> Option> { + pub const fn as_class_literal(self) -> Option> { match self { Type::ClassLiteral(class_type) => Some(class_type), _ => None, diff --git a/crates/ty_python_semantic/src/types/class.rs b/crates/ty_python_semantic/src/types/class.rs index 7b89f0828b..f85dd3ca45 100644 --- a/crates/ty_python_semantic/src/types/class.rs +++ b/crates/ty_python_semantic/src/types/class.rs @@ -1940,7 +1940,7 @@ impl<'db> ClassLiteral<'db> { #[salsa::tracked(cycle_initial=is_typed_dict_cycle_initial, heap_size=ruff_memory_usage::heap_size )] - pub(super) fn is_typed_dict(self, db: &'db dyn Db) -> bool { + pub fn is_typed_dict(self, db: &'db dyn Db) -> bool { if let Some(known) = self.known(db) { return known.is_typed_dict_subclass(); }