From 518d11b33f0f7acc48d9e79c15728dedd6e25239 Mon Sep 17 00:00:00 2001 From: Andrew Gallant Date: Thu, 4 Dec 2025 14:33:42 -0500 Subject: [PATCH] [ty] Add modules to auto-import This makes auto-import include modules in suggestions. In this initial implementation, we permit this to include submodules as well. This is in contrast to what we do in `import ...` completions. It's easy to change this behavior, but I think it'd be interesting to run with this for now to see how well it works. --- crates/ty_ide/src/all_symbols.rs | 74 +++++++++---- crates/ty_ide/src/completion.rs | 100 +++++++++++++++++- ...action_attribute_access_on_unimported.snap | 46 ++++++++ 3 files changed, 196 insertions(+), 24 deletions(-) diff --git a/crates/ty_ide/src/all_symbols.rs b/crates/ty_ide/src/all_symbols.rs index 210b116e57..aa7f9e02b7 100644 --- a/crates/ty_ide/src/all_symbols.rs +++ b/crates/ty_ide/src/all_symbols.rs @@ -59,12 +59,19 @@ pub fn all_symbols<'db>( continue; } s.spawn(move |_| { + if query.is_match_symbol_name(module.name(&*db)) { + results.lock().unwrap().push(AllSymbolInfo { + symbol: None, + module, + file, + }); + } for (_, symbol) in symbols_for_file_global_only(&*db, file).search(query) { // It seems like we could do better here than // locking `results` for every single symbol, // but this works pretty well as it is. results.lock().unwrap().push(AllSymbolInfo { - symbol: symbol.to_owned(), + symbol: Some(symbol.to_owned()), module, file, }); @@ -76,8 +83,16 @@ pub fn all_symbols<'db>( let mut results = results.into_inner().unwrap(); results.sort_by(|s1, s2| { - let key1 = (&s1.symbol.name, s1.file.path(db).as_str()); - let key2 = (&s2.symbol.name, s2.file.path(db).as_str()); + let key1 = ( + s1.name_in_file() + .unwrap_or_else(|| s1.module().name(db).as_str()), + s1.file.path(db).as_str(), + ); + let key2 = ( + s2.name_in_file() + .unwrap_or_else(|| s2.module().name(db).as_str()), + s2.file.path(db).as_str(), + ); key1.cmp(&key2) }); results @@ -88,7 +103,9 @@ pub fn all_symbols<'db>( #[derive(Debug, Clone, PartialEq, Eq)] pub struct AllSymbolInfo<'db> { /// The symbol information. - symbol: SymbolInfo<'static>, + /// + /// When absent, this implies the symbol is the module itself. + symbol: Option>, /// The module containing the symbol. module: Module<'db>, /// The file containing the symbol. @@ -99,9 +116,14 @@ pub struct AllSymbolInfo<'db> { } impl<'db> AllSymbolInfo<'db> { - /// Returns the name of this symbol. - pub fn name(&self) -> &str { - &self.symbol.name + /// Returns the name of this symbol as it exists in a file. + /// + /// When absent, there is no concrete symbol in a module + /// somewhere. Instead, this represents importing a module. + /// In this case, if the caller needs a symbol name, they + /// should use `AllSymbolInfo::module().name()`. + pub fn name_in_file(&self) -> Option<&str> { + self.symbol.as_ref().map(|symbol| &*symbol.name) } /// Returns the "kind" of this symbol. @@ -110,7 +132,10 @@ impl<'db> AllSymbolInfo<'db> { /// determined on a best effort basis. It may be imprecise /// in some cases, e.g., reporting a module as a variable. pub fn kind(&self) -> SymbolKind { - self.symbol.kind + self.symbol + .as_ref() + .map(|symbol| symbol.kind) + .unwrap_or(SymbolKind::Module) } /// Returns the module this symbol is exported from. @@ -208,25 +233,31 @@ ABCDEFGHIJKLMNOP = 'https://api.example.com' return "No symbols found".to_string(); } - self.render_diagnostics(symbols.into_iter().map(AllSymbolDiagnostic::new)) + self.render_diagnostics(symbols.into_iter().map(|symbol_info| AllSymbolDiagnostic { + db: &self.db, + symbol_info, + })) } } struct AllSymbolDiagnostic<'db> { + db: &'db dyn Db, symbol_info: AllSymbolInfo<'db>, } - impl<'db> AllSymbolDiagnostic<'db> { - fn new(symbol_info: AllSymbolInfo<'db>) -> Self { - Self { symbol_info } - } - } - impl IntoDiagnostic for AllSymbolDiagnostic<'_> { fn into_diagnostic(self) -> Diagnostic { - let symbol_kind_str = self.symbol_info.symbol.kind.to_string(); + let symbol_kind_str = self.symbol_info.kind().to_string(); - let info_text = format!("{} {}", symbol_kind_str, self.symbol_info.symbol.name); + let info_text = format!( + "{} {}", + symbol_kind_str, + self.symbol_info.name_in_file().unwrap_or_else(|| self + .symbol_info + .module() + .name(self.db) + .as_str()) + ); let sub = SubDiagnostic::new(SubDiagnosticSeverity::Info, info_text); @@ -235,9 +266,12 @@ ABCDEFGHIJKLMNOP = 'https://api.example.com' Severity::Info, "AllSymbolInfo".to_string(), ); - main.annotate(Annotation::primary( - Span::from(self.symbol_info.file).with_range(self.symbol_info.symbol.name_range), - )); + + let mut span = Span::from(self.symbol_info.file()); + if let Some(ref symbol) = self.symbol_info.symbol { + span = span.with_range(symbol.name_range); + } + main.annotate(Annotation::primary(span)); main.sub(sub); main diff --git a/crates/ty_ide/src/completion.rs b/crates/ty_ide/src/completion.rs index f96ac1500d..70505ac4c8 100644 --- a/crates/ty_ide/src/completion.rs +++ b/crates/ty_ide/src/completion.rs @@ -74,7 +74,7 @@ impl<'db> Completions<'db> { .into_iter() .filter_map(|item| { Some(ImportEdit { - label: format!("import {}.{}", item.module_name?, item.name), + label: format!("import {}", item.qualified?), edit: item.import?, }) }) @@ -160,6 +160,10 @@ impl<'db> Extend> for Completions<'db> { pub struct Completion<'db> { /// The label shown to the user for this suggestion. pub name: Name, + /// The fully qualified name, when available. + /// + /// This is only set when `module_name` is available. + pub qualified: Option, /// The text that should be inserted at the cursor /// when the completion is selected. /// @@ -225,6 +229,7 @@ impl<'db> Completion<'db> { let is_type_check_only = semantic.is_type_check_only(db); Completion { name: semantic.name, + qualified: None, insert: None, ty: semantic.ty, kind: None, @@ -306,6 +311,7 @@ impl<'db> Completion<'db> { fn keyword(name: &str) -> Self { Completion { name: name.into(), + qualified: None, insert: None, ty: None, kind: Some(CompletionKind::Keyword), @@ -321,6 +327,7 @@ impl<'db> Completion<'db> { fn value_keyword(name: &str, ty: Type<'db>) -> Completion<'db> { Completion { name: name.into(), + qualified: None, insert: None, ty: Some(ty), kind: Some(CompletionKind::Keyword), @@ -541,7 +548,18 @@ fn add_unimported_completions<'db>( continue; } - let request = create_import_request(symbol.module().name(db), symbol.name()); + let module_name = symbol.module().name(db); + let (name, qualified, request) = symbol + .name_in_file() + .map(|name| { + let qualified = format!("{module_name}.{name}"); + (name, qualified, create_import_request(module_name, name)) + }) + .unwrap_or_else(|| { + let name = module_name.as_str(); + let qualified = name.to_string(); + (name, qualified, ImportRequest::module(name)) + }); // FIXME: `all_symbols` doesn't account for wildcard imports. // Since we're looking at every module, this is probably // "fine," but it might mean that we import a symbol from the @@ -550,11 +568,12 @@ fn add_unimported_completions<'db>( // N.B. We use `add` here because `all_symbols` already // takes our query into account. completions.force_add(Completion { - name: ast::name::Name::new(symbol.name()), + name: ast::name::Name::new(name), + qualified: Some(ast::name::Name::new(qualified)), insert: Some(import_action.symbol_text().into()), ty: None, kind: symbol.kind().to_completion_kind(), - module_name: Some(symbol.module().name(db)), + module_name: Some(module_name), import: import_action.import().cloned(), builtin: false, // TODO: `is_type_check_only` requires inferring the type of the symbol @@ -6066,6 +6085,79 @@ ZQ "); } + #[test] + fn auto_import_includes_stdlib_modules_as_suggestions() { + let snapshot = CursorTest::builder() + .source( + "main.py", + r#" +multiprocess +"#, + ) + .completion_test_builder() + .auto_import() + .build() + .snapshot(); + assert_snapshot!(snapshot, @r" + multiprocessing + multiprocessing.connection + multiprocessing.context + multiprocessing.dummy + multiprocessing.dummy.connection + multiprocessing.forkserver + multiprocessing.heap + multiprocessing.managers + multiprocessing.pool + multiprocessing.popen_fork + multiprocessing.popen_forkserver + multiprocessing.popen_spawn_posix + multiprocessing.popen_spawn_win32 + multiprocessing.process + multiprocessing.queues + multiprocessing.reduction + multiprocessing.resource_sharer + multiprocessing.resource_tracker + multiprocessing.shared_memory + multiprocessing.sharedctypes + multiprocessing.spawn + multiprocessing.synchronize + multiprocessing.util + "); + } + + #[test] + fn auto_import_includes_first_party_modules_as_suggestions() { + let snapshot = CursorTest::builder() + .source( + "main.py", + r#" +zqzqzq +"#, + ) + .source("zqzqzqzqzq.py", "") + .completion_test_builder() + .auto_import() + .build() + .snapshot(); + assert_snapshot!(snapshot, @"zqzqzqzqzq"); + } + + #[test] + fn auto_import_includes_sub_modules_as_suggestions() { + let snapshot = CursorTest::builder() + .source( + "main.py", + r#" +collabc +"#, + ) + .completion_test_builder() + .auto_import() + .build() + .snapshot(); + assert_snapshot!(snapshot, @"collections.abc"); + } + /// A way to create a simple single-file (named `main.py`) completion test /// builder. /// diff --git a/crates/ty_server/tests/e2e/snapshots/e2e__code_actions__code_action_attribute_access_on_unimported.snap b/crates/ty_server/tests/e2e/snapshots/e2e__code_actions__code_action_attribute_access_on_unimported.snap index c82a14bc8e..3026696d8e 100644 --- a/crates/ty_server/tests/e2e/snapshots/e2e__code_actions__code_action_attribute_access_on_unimported.snap +++ b/crates/ty_server/tests/e2e/snapshots/e2e__code_actions__code_action_attribute_access_on_unimported.snap @@ -3,6 +3,52 @@ source: crates/ty_server/tests/e2e/code_actions.rs expression: code_actions --- [ + { + "title": "import typing", + "kind": "quickfix", + "diagnostics": [ + { + "range": { + "start": { + "line": 0, + "character": 3 + }, + "end": { + "line": 0, + "character": 9 + } + }, + "severity": 1, + "code": "unresolved-reference", + "codeDescription": { + "href": "https://ty.dev/rules#unresolved-reference" + }, + "source": "ty", + "message": "Name `typing` used when not defined", + "relatedInformation": [] + } + ], + "edit": { + "changes": { + "file:///src/foo.py": [ + { + "range": { + "start": { + "line": 0, + "character": 0 + }, + "end": { + "line": 0, + "character": 0 + } + }, + "newText": "import typing\n" + } + ] + } + }, + "isPreferred": true + }, { "title": "Ignore 'unresolved-reference' for this line", "kind": "quickfix",