mirrors/ruff
1
0
Fork 0
mirror of https://github.com/astral-sh/ruff.git synced 2025-08-04 18:58:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

[ty] Add completions for from module import <CURSOR> (#18830)

Browse source 
There were two main challenges in this PR.

The first was mostly just figuring out how to get the symbols
corresponding to `module`. It turns out that we do this in a couple
of places in ty already, but through different means. In one approach,
we use [`exported_names`]. In another approach, we get a `Type`
corresponding to the module. We take the latter approach here, which is
consistent with how we do completions elsewhere. (I looked into
factoring this logic out into its own function, but it ended up being
pretty constrained. e.g., There's only one other place where we want to
go from `ast::StmtImportFrom` to a module `Type`, and that code also
wants the module name.)

The second challenge was recognizing the `from module import <CURSOR>`
pattern in the code. I initially started with some fixed token patterns
to get a proof of concept working. But I ended up switching to mini
state machine over tokens. I looked at the parser for `StmtImportFrom`
to determine what kinds of tokens we can expect.

[`exported_names`]:
23a3b6ef23/crates/ty_python_semantic/src/semantic_index/re_exports.rs (L47)
This commit is contained in:
Andrew Gallant 2025-06-23 10:43:25 -04:00 • committed by GitHub
parent 9e9c4fe17b
commit a77db3da3f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 342 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

317
crates/ty_ide/src/completion.rs
View file

@ -26,6 +26,7 @@ pub fn completion(db: &dyn Db, file: File, offset: TextSize) -> Vec<Completion>
let model = ty_python_semantic::SemanticModel::new(db.upcast(), file);
let mut completions = match target {
CompletionTargetAst::ObjectDot { expr } => model.attribute_completions(expr),
CompletionTargetAst::ImportFrom { import, name } => model.import_completions(import, name),
CompletionTargetAst::Scoped { node } => model.scoped_completions(node),
};
completions.sort_by(|name1, name2| compare_suggestions(name1, name2));
@ -62,6 +63,12 @@ enum CompletionTargetTokens<'t> {
#[expect(dead_code)]
attribute: Option<&'t Token>,
},
/// A `from module import attribute` token form was found, where
/// `attribute` may be empty.
ImportFrom {
/// The module being imported from.
module: &'t Token,
},
/// A token was found under the cursor, but it didn't
/// match any of our anticipated token patterns.
Generic { token: &'t Token },
@ -102,6 +109,8 @@ impl<'t> CompletionTargetTokens<'t> {
object,
attribute: Some(attribute),
}
} else if let Some(module) = import_from_tokens(before) {
CompletionTargetTokens::ImportFrom { module }
} else if let Some([_]) = token_suffix_by_kinds(before, [TokenKind::Float]) {
// If we're writing a `float`, then we should
// specifically not offer completions. This wouldn't
@ -153,6 +162,15 @@ impl<'t> CompletionTargetTokens<'t> {
_ => None,
}
}
CompletionTargetTokens::ImportFrom { module, .. } => {
let covering_node = covering_node(parsed.syntax().into(), module.range())
.find_first(|node| node.is_stmt_import_from())
.ok()?;
let ast::AnyNodeRef::StmtImportFrom(import) = covering_node.node() else {
return None;
};
Some(CompletionTargetAst::ImportFrom { import, name: None })
}
CompletionTargetTokens::Generic { token } => {
let covering_node = covering_node(parsed.syntax().into(), token.range());
Some(CompletionTargetAst::Scoped {
@ -176,6 +194,15 @@ enum CompletionTargetAst<'t> {
/// A `object.attribute` scenario, where we want to
/// list attributes on `object` for completions.
ObjectDot { expr: &'t ast::ExprAttribute },
/// A `from module import attribute` scenario, where we want to
/// list attributes on `module` for completions.
ImportFrom {
/// The import statement.
import: &'t ast::StmtImportFrom,
/// An index into `import.names` if relevant. When this is
/// set, the index is guaranteed to be valid.
name: Option<usize>,
},
/// A scoped scenario, where we want to list all items available in
/// the most narrow scope containing the giving AST node.
Scoped { node: ast::AnyNodeRef<'t> },
@ -205,6 +232,97 @@ fn token_suffix_by_kinds<const N: usize>(
}))
}
/// Looks for the start of a `from module import <CURSOR>` statement.
///
/// If found, one arbitrary token forming `module` is returned.
fn import_from_tokens(tokens: &[Token]) -> Option<&Token> {
use TokenKind as TK;
/// The number of tokens we're willing to consume backwards from
/// the cursor's position until we give up looking for a `from
/// module import <CURSOR>` pattern. The state machine below has
/// lots of opportunities to bail way earlier than this, but if
/// there's, e.g., a long list of name tokens for something that
/// isn't an import, then we could end up doing a lot of wasted
/// work here. Probably humans aren't often working with single
/// import statements over 1,000 tokens long.
///
/// The other thing to consider here is that, by the time we get to
/// this point, ty has already done some work proportional to the
/// length of `tokens` anyway. The unit of work we do below is very
/// small.
const LIMIT: usize = 1_000;
/// A state used to "parse" the tokens preceding the user's cursor,
/// in reverse, to detect a "from import" statement.
enum S {
Start,
Names,
Module,
}
let mut state = S::Start;
let mut module_token: Option<&Token> = None;
// Move backward through the tokens until we get to
// the `from` token.
for token in tokens.iter().rev().take(LIMIT) {
state = match (state, token.kind()) {
// It's okay to pop off a newline token here initially,
// since it may occur when the name being imported is
// empty.
(S::Start, TK::Newline) => S::Names,
// Munch through tokens that can make up an alias.
// N.B. We could also consider taking any token here
// *except* some limited set of tokens (like `Newline`).
// That might work well if it turns out that listing
// all possible allowable tokens is too brittle.
(
S::Start | S::Names,
TK::Name
| TK::Comma
| TK::As
| TK::Case
| TK::Match
| TK::Type
| TK::Star
| TK::Lpar
| TK::Rpar
| TK::NonLogicalNewline
// It's not totally clear the conditions under
// which this occurs (I haven't read our tokenizer),
// but it appears in code like this, where this is
// the entire file contents:
//
// from sys import (
// abiflags,
// <CURSOR>
//
// It seems harmless to just allow this "unknown"
// token here to make the above work.
| TK::Unknown,
) => S::Names,
(S::Start | S::Names, TK::Import) => S::Module,
// Munch through tokens that can make up a module.
(
S::Module,
TK::Name | TK::Dot | TK::Ellipsis | TK::Case | TK::Match | TK::Type | TK::Unknown,
) => {
// It's okay if there are multiple module
// tokens here. Just taking the last one
// (which is the one appearing first in
// the source code) is fine. We only need
// this to find the corresponding AST node,
// so any of the tokens should work fine.
module_token = Some(token);
S::Module
}
(S::Module, TK::From) => return module_token,
_ => return None,
};
}
None
}
/// Order completions lexicographically, with these exceptions:
///
/// 1) A `_[^_]` prefix sorts last and
@ -1850,6 +1968,205 @@ def test_point(p2: Point):
test.assert_completions_include("orthogonal_direction");
}
#[test]
fn from_import1() {
let test = cursor_test(
"\
from sys import <CURSOR>
",
);
test.assert_completions_include("getsizeof");
}
#[test]
fn from_import2() {
let test = cursor_test(
"\
from sys import abiflags, <CURSOR>
",
);
test.assert_completions_include("getsizeof");
}
#[test]
fn from_import3() {
let test = cursor_test(
"\
from sys import <CURSOR>, abiflags
",
);
test.assert_completions_include("getsizeof");
}
#[test]
fn from_import4() {
let test = cursor_test(
"\
from sys import abiflags, \
<CURSOR>
",
);
test.assert_completions_include("getsizeof");
}
#[test]
fn from_import5() {
let test = cursor_test(
"\
from sys import abiflags as foo, <CURSOR>
",
);
test.assert_completions_include("getsizeof");
}
#[test]
fn from_import6() {
let test = cursor_test(
"\
from sys import abiflags as foo, g<CURSOR>
",
);
test.assert_completions_include("getsizeof");
}
#[test]
fn from_import7() {
let test = cursor_test(
"\
from sys import abiflags as foo, \
<CURSOR>
",
);
test.assert_completions_include("getsizeof");
}
#[test]
fn from_import8() {
let test = cursor_test(
"\
from sys import abiflags as foo, \
g<CURSOR>
",
);
test.assert_completions_include("getsizeof");
}
#[test]
fn from_import9() {
let test = cursor_test(
"\
from sys import (
abiflags,
<CURSOR>
",
);
test.assert_completions_include("getsizeof");
}
#[test]
fn from_import10() {
let test = cursor_test(
"\
from sys import (
abiflags,
<CURSOR>
)
",
);
test.assert_completions_include("getsizeof");
}
#[test]
fn from_import11() {
let test = cursor_test(
"\
from sys import (
<CURSOR>
)
",
);
test.assert_completions_include("getsizeof");
}
#[test]
fn from_import_unknown_in_module() {
let test = cursor_test(
"\
foo = 1
from ? import <CURSOR>
",
);
assert_snapshot!(test.completions(), @r"<No completions found>");
}
#[test]
fn from_import_unknown_in_import_names1() {
let test = cursor_test(
"\
from sys import ?, <CURSOR>
",
);
test.assert_completions_include("getsizeof");
}
#[test]
fn from_import_unknown_in_import_names2() {
let test = cursor_test(
"\
from sys import ??, <CURSOR>
",
);
test.assert_completions_include("getsizeof");
}
#[test]
fn from_import_unknown_in_import_names3() {
let test = cursor_test(
"\
from sys import ??, <CURSOR>, ??
",
);
test.assert_completions_include("getsizeof");
}
#[test]
fn import_submodule_not_attribute1() {
let test = cursor_test(
"\
import importlib
importlib.<CURSOR>
",
);
test.assert_completions_do_not_include("resources");
}
#[test]
fn import_submodule_not_attribute2() {
let test = cursor_test(
"\
import importlib.resources
importlib.<CURSOR>
",
);
// TODO: This is wrong. Completions should include
// `resources` here.
test.assert_completions_do_not_include("resources");
}
#[test]
fn import_submodule_not_attribute3() {
let test = cursor_test(
"\
import importlib
import importlib.resources
importlib.<CURSOR>
",
);
// TODO: This is wrong. Completions should include
// `resources` here.
test.assert_completions_do_not_include("resources");
}
#[test]
fn regression_test_issue_642() {
// Regression test for https://github.com/astral-sh/ty/issues/642