diff --git a/crates/ty_ide/src/completion.rs b/crates/ty_ide/src/completion.rs index 3d6a5fc78b..4c6e4e93a0 100644 --- a/crates/ty_ide/src/completion.rs +++ b/crates/ty_ide/src/completion.rs @@ -358,6 +358,9 @@ fn only_keyword_completion<'db>(tokens: &[Token], typed: Option<&str>) -> Option if is_import_from_incomplete(tokens, typed) { return Some(Completion::keyword("import")); } + if is_import_alias_incomplete(tokens, typed) { + return Some(Completion::keyword("as")); + } None } @@ -914,6 +917,59 @@ fn is_import_from_incomplete(tokens: &[Token], typed: Option<&str>) -> bool { false } +/// Detects `import ` statements with a potentially incomplete +/// `as` clause. +/// +/// Note that this works for `from import ` as well. +/// +/// If found, `true` is returned. +fn is_import_alias_incomplete(tokens: &[Token], typed: Option<&str>) -> bool { + 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 "import as" statement. + enum S { + Start, + As, + Name, + } + + if typed.is_none() { + return false; + } + + let mut state = S::Start; + for token in tokens.iter().rev().take(LIMIT) { + state = match (state, token.kind()) { + (S::Start, TK::Name | TK::Unknown | TK::As) => S::As, + (S::As, TK::Name | TK::Case | TK::Match | TK::Type | TK::Unknown) => S::Name, + ( + S::Name, + TK::Name + | TK::Dot + | TK::Ellipsis + | TK::Case + | TK::Match + | TK::Type + | TK::Unknown + | TK::Comma + | TK::As + | TK::Newline + | TK::NonLogicalNewline + | TK::Lpar + | TK::Rpar, + ) => S::Name, + // Once we reach the `import` token we know we're in + // `import name `. + (S::Name, TK::Import) => return true, + _ => return false, + }; + } + false +} + /// Looks for the text typed immediately before the cursor offset /// given. /// @@ -4690,6 +4746,68 @@ from collections import defaultdict as f ); } + #[test] + fn import_missing_alias_suggests_as() { + let builder = completion_test_builder( + "\ +import collections a + ", + ); + assert_snapshot!(builder.build().snapshot(), @"as"); + } + + #[test] + fn import_dotted_module_missing_alias_suggests_as() { + let builder = completion_test_builder( + "\ +import collections.abc a + ", + ); + assert_snapshot!(builder.build().snapshot(), @"as"); + } + + #[test] + fn import_multiple_modules_missing_alias_suggests_as() { + let builder = completion_test_builder( + "\ +import collections.abc as c, typing a + ", + ); + assert_snapshot!(builder.build().snapshot(), @"as"); + } + + #[test] + fn from_import_missing_alias_suggests_as() { + let builder = completion_test_builder( + "\ +from collections.abc import Mapping a + ", + ); + assert_snapshot!(builder.build().snapshot(), @"as"); + } + + #[test] + fn from_import_parenthesized_missing_alias_suggests_as() { + let builder = completion_test_builder( + "\ +from typing import ( + NamedTuple a +) + ", + ); + assert_snapshot!(builder.build().snapshot(), @"as"); + } + + #[test] + fn from_relative_import_missing_alias_suggests_as() { + let builder = completion_test_builder( + "\ +from ...foo import bar a + ", + ); + assert_snapshot!(builder.build().snapshot(), @"as"); + } + #[test] fn no_completions_in_with_alias() { let builder = completion_test_builder(