[ty] Add support for module-only import requests

The existing importer functionality always required
an import request with a module and a member in that
module. But we want to be able to insert import statements
for a module itself and not any members in the module.

This is basically changing `member: &str` to an
`Option<&str>` and fixing the fallout in a way that
makes sense for module-only imports.
This commit is contained in:
Andrew Gallant 2025-12-04 14:33:21 -05:00 committed by Andrew Gallant
parent 3c2cf49f60
commit da94b99248
1 changed files with 134 additions and 16 deletions

View File

@ -145,7 +145,7 @@ impl<'a> Importer<'a> {
members: &MembersInScope, members: &MembersInScope,
) -> ImportAction { ) -> ImportAction {
let request = request.avoid_conflicts(self.db, self.file, members); let request = request.avoid_conflicts(self.db, self.file, members);
let mut symbol_text: Box<str> = request.member.into(); let mut symbol_text: Box<str> = request.member.unwrap_or(request.module).into();
let Some(response) = self.find(&request, members.at) else { let Some(response) = self.find(&request, members.at) else {
let insertion = if let Some(future) = self.find_last_future_import(members.at) { let insertion = if let Some(future) = self.find_last_future_import(members.at) {
Insertion::end_of_statement(future.stmt, self.source, self.stylist) Insertion::end_of_statement(future.stmt, self.source, self.stylist)
@ -157,14 +157,27 @@ impl<'a> Importer<'a> {
Insertion::start_of_file(self.parsed.suite(), self.source, self.stylist, range) Insertion::start_of_file(self.parsed.suite(), self.source, self.stylist, range)
}; };
let import = insertion.into_edit(&request.to_string()); let import = insertion.into_edit(&request.to_string());
if matches!(request.style, ImportStyle::Import) { if let Some(member) = request.member
symbol_text = format!("{}.{}", request.module, request.member).into(); && matches!(request.style, ImportStyle::Import)
{
symbol_text = format!("{}.{}", request.module, member).into();
} }
return ImportAction { return ImportAction {
import: Some(import), import: Some(import),
symbol_text, symbol_text,
}; };
}; };
// When we just have a request to import a module (and not
// any members from that module), then the only way we can be
// here is if we found a pre-existing import that definitively
// satisfies the request. So we're done.
let Some(member) = request.member else {
return ImportAction {
import: None,
symbol_text,
};
};
match response.kind { match response.kind {
ImportResponseKind::Unqualified { ast, alias } => { ImportResponseKind::Unqualified { ast, alias } => {
let member = alias.asname.as_ref().unwrap_or(&alias.name).as_str(); let member = alias.asname.as_ref().unwrap_or(&alias.name).as_str();
@ -189,13 +202,10 @@ impl<'a> Importer<'a> {
let import = if let Some(insertion) = let import = if let Some(insertion) =
Insertion::existing_import(response.import.stmt, self.tokens) Insertion::existing_import(response.import.stmt, self.tokens)
{ {
insertion.into_edit(request.member) insertion.into_edit(member)
} else { } else {
Insertion::end_of_statement(response.import.stmt, self.source, self.stylist) Insertion::end_of_statement(response.import.stmt, self.source, self.stylist)
.into_edit(&format!( .into_edit(&format!("from {} import {member}", request.module))
"from {} import {}",
request.module, request.member
))
}; };
ImportAction { ImportAction {
import: Some(import), import: Some(import),
@ -481,6 +491,17 @@ impl<'ast> AstImportKind<'ast> {
Some(ImportResponseKind::Qualified { ast, alias }) Some(ImportResponseKind::Qualified { ast, alias })
} }
AstImportKind::ImportFrom(ast) => { AstImportKind::ImportFrom(ast) => {
// If the request is for a module itself, then we
// assume that it can never be satisfies by a
// `from ... import ...` statement. For example, a
// `request for collections.abc` needs an
// `import collections.abc`. Now, there could be a
// `from collections import abc`, and we could
// plausibly consider that a match and return a
// symbol text of `abc`. But it's not clear if that's
// the right choice or not.
let member = request.member?;
if request.force_style && !matches!(request.style, ImportStyle::ImportFrom) { if request.force_style && !matches!(request.style, ImportStyle::ImportFrom) {
return None; return None;
} }
@ -492,9 +513,7 @@ impl<'ast> AstImportKind<'ast> {
let kind = ast let kind = ast
.names .names
.iter() .iter()
.find(|alias| { .find(|alias| alias.name.as_str() == "*" || alias.name.as_str() == member)
alias.name.as_str() == "*" || alias.name.as_str() == request.member
})
.map(|alias| ImportResponseKind::Unqualified { ast, alias }) .map(|alias| ImportResponseKind::Unqualified { ast, alias })
.unwrap_or_else(|| ImportResponseKind::Partial(ast)); .unwrap_or_else(|| ImportResponseKind::Partial(ast));
Some(kind) Some(kind)
@ -510,7 +529,10 @@ pub(crate) struct ImportRequest<'a> {
/// `foo`, in `from foo import bar`). /// `foo`, in `from foo import bar`).
module: &'a str, module: &'a str,
/// The member to import (e.g., `bar`, in `from foo import bar`). /// The member to import (e.g., `bar`, in `from foo import bar`).
member: &'a str, ///
/// When `member` is absent, then this request reflects an import
/// of the module itself. i.e., `import module`.
member: Option<&'a str>,
/// The preferred style to use when importing the symbol (e.g., /// The preferred style to use when importing the symbol (e.g.,
/// `import foo` or `from foo import bar`). /// `import foo` or `from foo import bar`).
/// ///
@ -532,7 +554,7 @@ impl<'a> ImportRequest<'a> {
pub(crate) fn import(module: &'a str, member: &'a str) -> Self { pub(crate) fn import(module: &'a str, member: &'a str) -> Self {
Self { Self {
module, module,
member, member: Some(member),
style: ImportStyle::Import, style: ImportStyle::Import,
force_style: false, force_style: false,
} }
@ -545,12 +567,26 @@ impl<'a> ImportRequest<'a> {
pub(crate) fn import_from(module: &'a str, member: &'a str) -> Self { pub(crate) fn import_from(module: &'a str, member: &'a str) -> Self {
Self { Self {
module, module,
member, member: Some(member),
style: ImportStyle::ImportFrom, style: ImportStyle::ImportFrom,
force_style: false, force_style: false,
} }
} }
/// Create a new [`ImportRequest`] for bringing the given module
/// into scope.
///
/// This is for just importing the module itself, always via an
/// `import` statement.
pub(crate) fn module(module: &'a str) -> Self {
Self {
module,
member: None,
style: ImportStyle::Import,
force_style: false,
}
}
/// Causes this request to become a command. This will force the /// Causes this request to become a command. This will force the
/// requested import style, even if another style would be more /// requested import style, even if another style would be more
/// appropriate generally. /// appropriate generally.
@ -565,7 +601,13 @@ impl<'a> ImportRequest<'a> {
/// of an import conflict are minimized (although not always reduced /// of an import conflict are minimized (although not always reduced
/// to zero). /// to zero).
fn avoid_conflicts(self, db: &dyn Db, importing_file: File, members: &MembersInScope) -> Self { fn avoid_conflicts(self, db: &dyn Db, importing_file: File, members: &MembersInScope) -> Self {
match (members.map.get(self.module), members.map.get(self.member)) { let Some(member) = self.member else {
return Self {
style: ImportStyle::Import,
..self
};
};
match (members.map.get(self.module), members.map.get(member)) {
// Neither symbol exists, so we can just proceed as // Neither symbol exists, so we can just proceed as
// normal. // normal.
(None, None) => self, (None, None) => self,
@ -630,7 +672,10 @@ impl std::fmt::Display for ImportRequest<'_> {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
match self.style { match self.style {
ImportStyle::Import => write!(f, "import {}", self.module), ImportStyle::Import => write!(f, "import {}", self.module),
ImportStyle::ImportFrom => write!(f, "from {} import {}", self.module, self.member), ImportStyle::ImportFrom => match self.member {
None => write!(f, "import {}", self.module),
Some(member) => write!(f, "from {} import {member}", self.module),
},
} }
} }
} }
@ -843,6 +888,10 @@ mod tests {
self.add(ImportRequest::import_from(module, member)) self.add(ImportRequest::import_from(module, member))
} }
fn module(&self, module: &str) -> String {
self.add(ImportRequest::module(module))
}
fn add(&self, request: ImportRequest<'_>) -> String { fn add(&self, request: ImportRequest<'_>) -> String {
let node = covering_node( let node = covering_node(
self.cursor.parsed.syntax().into(), self.cursor.parsed.syntax().into(),
@ -2156,4 +2205,73 @@ except ImportError:
(bar.MAGIC) (bar.MAGIC)
"); ");
} }
#[test]
fn import_module_blank() {
let test = cursor_test(
"\
<CURSOR>
",
);
assert_snapshot!(
test.module("collections"), @r"
import collections
collections
");
}
#[test]
fn import_module_exists() {
let test = cursor_test(
"\
import collections
<CURSOR>
",
);
assert_snapshot!(
test.module("collections"), @r"
import collections
collections
");
}
#[test]
fn import_module_from_exists() {
let test = cursor_test(
"\
from collections import defaultdict
<CURSOR>
",
);
assert_snapshot!(
test.module("collections"), @r"
import collections
from collections import defaultdict
collections
");
}
// This test is working as intended. That is,
// `abc` is already in scope, so requesting an
// import for `collections.abc` could feasibly
// reuse the import and rewrite the symbol text
// to just `abc`. But for now it seems better
// to respect what has been written and add the
// `import collections.abc`. This behavior could
// plausibly be changed.
#[test]
fn import_module_from_via_member_exists() {
let test = cursor_test(
"\
from collections import abc
<CURSOR>
",
);
assert_snapshot!(
test.module("collections.abc"), @r"
import collections.abc
from collections import abc
collections.abc
");
}
} }