[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.
This commit is contained in:
Andrew Gallant 2025-12-04 14:33:42 -05:00 committed by Andrew Gallant
parent da94b99248
commit 518d11b33f
3 changed files with 196 additions and 24 deletions

View File

@ -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<SymbolInfo<'static>>,
/// 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

View File

@ -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<Completion<'db>> 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<Name>,
/// 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<CURSOR>
");
}
#[test]
fn auto_import_includes_stdlib_modules_as_suggestions() {
let snapshot = CursorTest::builder()
.source(
"main.py",
r#"
multiprocess<CURSOR>
"#,
)
.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<CURSOR>
"#,
)
.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<CURSOR>
"#,
)
.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.
///

View File

@ -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://<temp_dir>/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",