[ty] provide import completion when in from <name> <name> statement (#21291)

<!--
Thank you for contributing to Ruff/ty! To help us out with reviewing,
please consider the following:

- Does this pull request include a summary of the change? (See below.)
- Does this pull request include a descriptive title? (Please prefix
with `[ty]` for ty pull
  requests.)
- Does this pull request include references to any relevant issues?
-->

## Summary

Resolves https://github.com/astral-sh/ty/issues/1494

## Test Plan

Add a test showing if we are in `from <name> <name> ` we provide the
keyword completion "import"
This commit is contained in:
Matthew Mckee 2025-11-10 16:59:45 +00:00 committed by GitHub
parent 4821c050ef
commit 1d188476b6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -160,6 +160,20 @@ impl<'db> Completion<'db> {
.and_then(|ty| imp(db, ty, &CompletionKindVisitor::default()))
})
}
fn keyword(name: &str) -> Self {
Completion {
name: name.into(),
insert: None,
ty: None,
kind: Some(CompletionKind::Keyword),
module_name: None,
import: None,
builtin: false,
is_type_check_only: false,
documentation: None,
}
}
}
/// The "kind" of a completion.
@ -212,14 +226,16 @@ pub fn completion<'db>(
offset: TextSize,
) -> Vec<Completion<'db>> {
let parsed = parsed_module(db, file).load(db);
let tokens = tokens_start_before(parsed.tokens(), offset);
let typed = find_typed_text(db, file, &parsed, offset);
if is_in_comment(tokens) || is_in_string(tokens) || is_in_definition_place(db, tokens, file) {
if is_in_no_completions_place(db, tokens, file) {
return vec![];
}
if let Some(completions) = only_keyword_completion(tokens, typed.as_deref()) {
return vec![completions];
}
let typed = find_typed_text(db, file, &parsed, offset);
let typed_query = typed
.as_deref()
.map(QueryPattern::new)
@ -309,6 +325,17 @@ fn add_keyword_value_completions<'db>(
}
}
/// When the tokens indicate that the last token should be precisely one
/// possible keyword, we provide a single completion for it.
///
/// `typed` should be the text that we think the user has typed so far.
fn only_keyword_completion<'db>(tokens: &[Token], typed: Option<&str>) -> Option<Completion<'db>> {
if is_import_from_incomplete(tokens, typed) {
return Some(Completion::keyword("import"));
}
None
}
/// Adds completions not in scope.
///
/// `scoped` should be information about the identified scope
@ -801,6 +828,67 @@ fn import_tokens(tokens: &[Token]) -> Option<(&Token, &Token)> {
None
}
/// Looks for the start of a `from module <CURSOR>` statement.
///
/// If found, `true` is returned.
///
/// `typed` should be the text that we think the user has typed so far.
fn is_import_from_incomplete(tokens: &[Token], typed: Option<&str>) -> bool {
// N.B. The implementation here is very similar to
// `from_import_tokens`. The main difference is that
// we're just looking for whether we should suggest
// the `import` keyword. So this is a little simpler.
use TokenKind as TK;
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,
ImportKeyword,
ModulePossiblyDotted,
ModuleOnlyDotted,
}
let mut state = S::Start;
if typed.is_none() {
state = S::ImportKeyword;
}
// 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()) {
// Match an incomplete `import` keyword.
//
// It's okay to pop off a newline token here initially,
// since it may occur before the user starts typing
// `import` but after the module name.
(S::Start, TK::Newline | TK::Name | TK::Import) => S::ImportKeyword,
// We are a bit more careful with how we parse the module
// here than in `from_import_tokens`. In particular, we
// want to make sure we don't incorrectly suggest `import`
// for `from os.i<CURSOR>`. If we aren't careful, then
// `i` could be considered an incomplete `import` keyword
// and `os.` is the module. But of course, ending with a
// `.` (unless the entire module is dots) is invalid.
(S::ImportKeyword, TK::Dot | TK::Ellipsis) => S::ModuleOnlyDotted,
(S::ImportKeyword, TK::Name | TK::Case | TK::Match | TK::Type | TK::Unknown) => {
S::ModulePossiblyDotted
}
(S::ModuleOnlyDotted, TK::Dot | TK::Ellipsis) => S::ModuleOnlyDotted,
(
S::ModulePossiblyDotted,
TK::Name | TK::Dot | TK::Ellipsis | TK::Case | TK::Match | TK::Type | TK::Unknown,
) => S::ModulePossiblyDotted,
(S::ModulePossiblyDotted | S::ModuleOnlyDotted, TK::From) => return true,
_ => return false,
};
}
false
}
/// Looks for the text typed immediately before the cursor offset
/// given.
///
@ -815,7 +903,10 @@ fn find_typed_text(
let source = source_text(db, file);
let tokens = tokens_start_before(parsed.tokens(), offset);
let last = tokens.last()?;
if !matches!(last.kind(), TokenKind::Name) {
// It's odd to include `TokenKind::Import` here, but it
// indicates that the user has typed `import`. This is
// useful to know in some contexts.
if !matches!(last.kind(), TokenKind::Name | TokenKind::Import) {
return None;
}
// This one's weird, but if the cursor is beyond
@ -830,6 +921,11 @@ fn find_typed_text(
Some(source[last.range()].to_string())
}
/// Whether the last token is in a place where we should not provide completions.
fn is_in_no_completions_place(db: &dyn Db, tokens: &[Token], file: File) -> bool {
is_in_comment(tokens) || is_in_string(tokens) || is_in_definition_place(db, tokens, file)
}
/// Whether the last token is within a comment or not.
fn is_in_comment(tokens: &[Token]) -> bool {
tokens.last().is_some_and(|t| t.kind().is_comment())
@ -4216,6 +4312,138 @@ type <CURSOR>
");
}
#[test]
fn from_import_i_suggests_import() {
let builder = completion_test_builder("from typing i<CURSOR>");
assert_snapshot!(builder.build().snapshot(), @"import");
}
#[test]
fn from_import_import_suggests_nothing() {
let builder = completion_test_builder("from typing import<CURSOR>");
assert_snapshot!(builder.build().snapshot(), @"import");
}
#[test]
fn from_import_importt_suggests_import() {
let builder = completion_test_builder("from typing importt<CURSOR>");
assert_snapshot!(builder.build().snapshot(), @"import");
}
#[test]
fn from_import_space_suggests_import() {
let builder = completion_test_builder("from typing <CURSOR>");
assert_snapshot!(builder.build().snapshot(), @"import");
}
#[test]
fn from_import_no_space_not_suggests_import() {
let builder = completion_test_builder("from typing<CURSOR>");
assert_snapshot!(builder.build().snapshot(), @r"
typing
typing_extensions
");
}
#[test]
fn from_import_two_imports_suggests_import() {
let builder = completion_test_builder(
"from collections.abc import Sequence
from typing i<CURSOR>",
);
assert_snapshot!(builder.build().snapshot(), @"import");
}
/// The following behaviour may not be reflected in editors, since LSP
/// clients may do their own filtering of completion suggestions.
#[test]
fn from_import_random_name_suggests_import() {
let builder = completion_test_builder("from typing aa<CURSOR>");
assert_snapshot!(builder.build().snapshot(), @"import");
}
#[test]
fn from_import_dotted_name_suggests_import() {
let builder = completion_test_builder("from collections.abc i<CURSOR>");
assert_snapshot!(builder.build().snapshot(), @"import");
}
#[test]
fn from_import_relative_import_suggests_import() {
let builder = CursorTest::builder()
.source("main.py", "from .foo i<CURSOR>")
.source("foo.py", "")
.completion_test_builder();
assert_snapshot!(builder.build().snapshot(), @"import");
}
#[test]
fn from_import_dotted_name_relative_import_suggests_import() {
let builder = CursorTest::builder()
.source("main.py", "from .foo.bar i<CURSOR>")
.source("foo/bar.py", "")
.completion_test_builder();
assert_snapshot!(builder.build().snapshot(), @"import");
}
#[test]
fn from_import_nested_dotted_name_relative_import_suggests_import() {
let builder = CursorTest::builder()
.source("src/main.py", "from ..foo i<CURSOR>")
.source("foo.py", "")
.completion_test_builder();
assert_snapshot!(builder.build().snapshot(), @"import");
}
#[test]
fn from_import_nested_very_dotted_name_relative_import_suggests_import() {
let builder = CursorTest::builder()
// N.B. the `...` tokenizes as `TokenKind::Ellipsis`
.source("src/main.py", "from ...foo i<CURSOR>")
.source("foo.py", "")
.completion_test_builder();
assert_snapshot!(builder.build().snapshot(), @"import");
}
#[test]
fn from_import_only_dot() {
let builder = CursorTest::builder()
.source(
"main.py",
"
import_zqzqzq = 1
from .<CURSOR>
",
)
.completion_test_builder();
assert_snapshot!(builder.build().snapshot(), @"import");
}
#[test]
fn from_import_only_dot_incomplete() {
let builder = CursorTest::builder()
.source(
"main.py",
"
import_zqzqzq = 1
from .imp<CURSOR>
",
)
.completion_test_builder();
assert_snapshot!(builder.build().snapshot(), @"import");
}
#[test]
fn from_import_incomplete() {
let builder = completion_test_builder(
"from collections.abc i
ZQZQZQ = 1
ZQ<CURSOR>",
);
assert_snapshot!(builder.build().snapshot(), @"ZQZQZQ");
}
/// A way to create a simple single-file (named `main.py`) completion test
/// builder.
///