From 553e5686248e86b9c781af77c714adc229bf8cc5 Mon Sep 17 00:00:00 2001 From: Andrew Gallant Date: Tue, 18 Nov 2025 10:52:01 -0500 Subject: [PATCH] [ty] Refactor detection of import statements for completions This commit essentially does away of all our old heuristic and piecemeal code for detecting different kinds of import statements. Instead, we offer one single state machine that does everything. This on its own fixes a few bugs. For example, `import collections.abc, unico` would previously offer global scope completions instead of module completions. For the most part though, this commit is a refactoring that preserves parity. In the next commit, we'll add support for completions on relative imports. --- crates/ty_ide/src/completion.rs | 1015 ++++++++++------- .../ty_python_semantic/src/semantic_model.rs | 69 +- 2 files changed, 600 insertions(+), 484 deletions(-) diff --git a/crates/ty_ide/src/completion.rs b/crates/ty_ide/src/completion.rs index 8563258571..c3ba8cf69d 100644 --- a/crates/ty_ide/src/completion.rs +++ b/crates/ty_ide/src/completion.rs @@ -332,55 +332,38 @@ pub fn completion<'db>( if is_in_no_completions_place(db, file, tokens, typed.as_deref()) { return vec![]; } - if let Some(completions) = only_keyword_completion(tokens, typed.as_deref()) { - return vec![completions]; - } - - let Some(target_token) = CompletionTargetTokens::find(&parsed, offset) else { - return vec![]; - }; - let Some(target) = target_token.ast(&parsed, offset) else { - return vec![]; - }; - - let model = SemanticModel::new(db, file); - let (semantic_completions, scoped) = match target { - CompletionTargetAst::ObjectDot { expr } => (model.attribute_completions(expr), None), - CompletionTargetAst::ObjectDotInImport { import, name } => { - (model.import_submodule_completions(import, name), None) - } - CompletionTargetAst::ObjectDotInImportFrom { import } => { - (model.from_import_submodule_completions(import), None) - } - CompletionTargetAst::ImportFrom { import, name } => { - (model.from_import_completions(import, name), None) - } - CompletionTargetAst::Import { .. } | CompletionTargetAst::ImportViaFrom { .. } => { - (model.import_completions(), None) - } - CompletionTargetAst::Scoped(scoped) => { - (model.scoped_completions(scoped.node), Some(scoped)) - } - }; let mut completions = Completions::new(db, typed.as_deref()); - completions.extend(semantic_completions); - if scoped.is_some() { - add_keyword_completions(db, &mut completions); - } - if settings.auto_import { - if let Some(scoped) = scoped { - add_unimported_completions( - db, - file, - &parsed, - scoped, - typed.as_deref(), - &mut completions, - ); + if let Some(import) = ImportStatement::detect(db, file, &parsed, tokens, typed.as_deref()) { + import.add_completions(db, file, &mut completions); + } else { + let Some(target_token) = CompletionTargetTokens::find(&parsed, offset) else { + return vec![]; + }; + let Some(target) = target_token.ast(&parsed, offset) else { + return vec![]; + }; + + let model = SemanticModel::new(db, file); + let (semantic_completions, scoped) = match target { + CompletionTargetAst::ObjectDot { expr } => (model.attribute_completions(expr), None), + CompletionTargetAst::Scoped(scoped) => { + (model.scoped_completions(scoped.node), Some(scoped)) + } + }; + + completions.extend(semantic_completions); + if scoped.is_some() { + add_keyword_completions(db, &mut completions); + } + if settings.auto_import { + if let Some(scoped) = scoped { + add_unimported_completions(db, file, &parsed, scoped, &mut completions); + } } } + completions.into_completions() } @@ -416,20 +399,6 @@ fn add_keyword_completions<'db>(db: &'db dyn Db, completions: &mut Completions<' } } -/// 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> { - 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 -} - /// Adds completions not in scope. /// /// `scoped` should be information about the identified scope @@ -442,7 +411,6 @@ fn add_unimported_completions<'db>( file: File, parsed: &ParsedModuleRef, scoped: ScopedTarget<'_>, - typed: Option<&str>, completions: &mut Completions<'db>, ) { // This is redundant since `all_symbols` will also bail @@ -470,7 +438,9 @@ fn add_unimported_completions<'db>( // "fine," but it might mean that we import a symbol from the // "wrong" module. let import_action = importer.import(request, &members); - completions.add(Completion { + // 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.symbol.name), insert: Some(import_action.symbol_text().into()), ty: None, @@ -509,27 +479,8 @@ enum CompletionTargetTokens<'t> { /// completions for. But we could use it for other things, /// like only returning completions that start with a prefix /// corresponding to this token. - 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 `import module` token form was found, where `module` may be - /// empty. - Import { - /// The token corresponding to the `import` keyword. - import: &'t Token, - /// The token closest to the cursor. - /// - /// This is currently unused, but we should use this - /// eventually to remove completions that aren't a - /// prefix of what has already been typed. (We are - /// currently relying on the LSP client to do this.) #[expect(dead_code)] - module: &'t Token, + attribute: Option<&'t Token>, }, /// A token was found under the cursor, but it didn't /// match any of our anticipated token patterns. @@ -571,10 +522,6 @@ impl<'t> CompletionTargetTokens<'t> { object, attribute: Some(attribute), } - } else if let Some(module) = import_from_tokens(before) { - CompletionTargetTokens::ImportFrom { module } - } else if let Some((import, module)) = import_tokens(before) { - CompletionTargetTokens::Import { import, 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 @@ -610,7 +557,7 @@ impl<'t> CompletionTargetTokens<'t> { offset: TextSize, ) -> Option> { match *self { - CompletionTargetTokens::PossibleObjectDot { object, attribute } => { + CompletionTargetTokens::PossibleObjectDot { object, .. } => { let covering_node = covering_node(parsed.syntax().into(), object.range()) .find_last(|node| { // We require that the end of the node range not @@ -636,44 +583,6 @@ impl<'t> CompletionTargetTokens<'t> { ast::AnyNodeRef::ExprAttribute(expr) => { Some(CompletionTargetAst::ObjectDot { expr }) } - ast::AnyNodeRef::StmtImport(import) => { - let range = attribute - .map(Ranged::range) - .unwrap_or_else(|| object.range()); - // Find the name that overlaps with the - // token we identified for the attribute. - let name = import - .names - .iter() - .position(|alias| alias.range().contains_range(range))?; - Some(CompletionTargetAst::ObjectDotInImport { import, name }) - } - ast::AnyNodeRef::StmtImportFrom(import) => { - Some(CompletionTargetAst::ObjectDotInImportFrom { import }) - } - _ => 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::Import { import, .. } => { - let covering_node = covering_node(parsed.syntax().into(), import.range()) - .find_first(|node| node.is_stmt_import() || node.is_stmt_import_from()) - .ok()?; - match covering_node.node() { - ast::AnyNodeRef::StmtImport(import) => { - Some(CompletionTargetAst::Import { import, name: None }) - } - ast::AnyNodeRef::StmtImportFrom(import) => { - Some(CompletionTargetAst::ImportViaFrom { import }) - } _ => None, } } @@ -698,45 +607,6 @@ enum CompletionTargetAst<'t> { /// A `object.attribute` scenario, where we want to /// list attributes on `object` for completions. ObjectDot { expr: &'t ast::ExprAttribute }, - /// A `import module.submodule` scenario, where we only want to - /// list submodules for completions. - ObjectDotInImport { - /// The import statement. - import: &'t ast::StmtImport, - /// An index into `import.names`. The index is guaranteed to be - /// valid. - name: usize, - }, - /// A `from module.submodule` scenario, where we only want to list - /// submodules for completions. - ObjectDotInImportFrom { import: &'t ast::StmtImportFrom }, - /// 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, - }, - /// A `import module` scenario, where we want to - /// list available modules for completions. - Import { - /// The import statement. - #[expect(dead_code)] - import: &'t ast::StmtImport, - /// An index into `import.names` if relevant. When this is - /// set, the index is guaranteed to be valid. - #[expect(dead_code)] - name: Option, - }, - /// A `from module` scenario, where we want to - /// list available modules for completions. - ImportViaFrom { - /// The import statement. - #[expect(dead_code)] - import: &'t ast::StmtImportFrom, - }, /// A scoped scenario, where we want to list all items available in /// the most narrow scope containing the giving AST node. Scoped(ScopedTarget<'t>), @@ -749,6 +619,528 @@ struct ScopedTarget<'t> { node: ast::AnyNodeRef<'t>, } +/// A representation of the completion context for a possibly incomplete import +/// statement. +#[derive(Clone, Debug)] +enum ImportStatement<'a> { + FromImport(FromImport<'a>), + Import(Import<'a>), + Incomplete(IncompleteImport), +} + +/// A representation of the completion context for a possibly incomplete +/// `from ... import ...` statement. +#[derive(Clone, Debug)] +struct FromImport<'a> { + ast: &'a ast::StmtImportFrom, + kind: FromImportKind, +} + +/// The kind of completions to offer for a `from import` statement. +/// +/// This is either something like `from col`, where we should +/// offer module completions, or `from collections.`, where +/// we should offer submodule completions or +/// `from collections import default` where we should offer +/// submodule/attribute completions. +#[derive(Clone, Debug)] +enum FromImportKind { + Module, + Submodule { parent: ModuleName }, + Attribute, +} + +/// A representation of the completion context for a possibly incomplete +/// `import ...` statement. +#[derive(Clone, Debug)] +struct Import<'a> { + #[expect(dead_code)] + ast: &'a ast::StmtImport, + kind: ImportKind, +} + +/// The kind of completions to offer for an `import` statement. +/// +/// This is either something like `import col`, where we should +/// offer module completions, or `import collections.`, where +/// we should offer submodule completions. +#[derive(Clone, Debug)] +enum ImportKind { + Module, + Submodule { parent: ModuleName }, +} + +/// Occurs when we detect that an import statement +/// is likely incomplete by virtue of a missing or +/// in-progress `as` or `import` keyword. +#[derive(Clone, Debug)] +enum IncompleteImport { + As, + Import, +} + +impl<'a> ImportStatement<'a> { + /// The number of tokens we're willing to consume backwards from + /// the cursor's position until we give up looking for an import + /// statement. 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; + + /// Attempts to detect an import statement in reverse starting at + /// the end of `tokens`. That is, `tokens` should correspond to the + /// sequence of tokens up to the end user's cursor. `typed` should + /// correspond to the text the user has typed, which is usually, + /// but not always, the text corresponding to the last token in + /// `tokens`. + fn detect( + db: &'a dyn Db, + file: File, + parsed: &'a ParsedModuleRef, + tokens: &'a [Token], + typed: Option<&str>, + ) -> Option> { + use TokenKind as TK; + + // This state machine moves backwards through the token stream, + // starting at where the user's cursor is and ending when + // either a `from` token is found, or a token that cannot + // possibly appear in an import statement at a particular + // position is found. + // + // To understand this state machine, it's recommended to become + // familiar with the grammar for Python import statements: + // https://docs.python.org/3/reference/grammar.html + + /// The current state of the parser below. + #[derive(Clone, Copy, Debug)] + enum S { + /// Our initial state. + Start, + /// We just saw an `import` token. + Import, + /// We just saw a first "name" token. That is, + /// a name-like token that appears just before + /// the end user's cursor. + /// + /// This isn't just limited to `TokenKind::Name`. + /// This also includes keywords and things like + /// "unknown" tokens that can stand in for names + /// at times. + FirstName, + /// A name seen immediately after the first name. This + /// indicates we have an incomplete import statement like + /// `import foo a` or `from foo imp`. But + /// we mush on. + AdjacentName, + + /// A state where we expect to see the start of or + /// continuation of a list of names following `import`. + /// In the [grammar], this is either `dotted_as_names` + /// or `import_from_as_names`. + /// + /// [grammar]: https://docs.python.org/3/reference/grammar.html + NameList, + /// Occurs after seeing a name-like token at the end + /// of a name list. This could be an alias, a dotted + /// name or a non-dotted name. + NameListNameOrAlias, + /// Occurs when we've seen an `as` in a list of names. + As, + /// Occurs when we see a name-like token after an `as` + /// keyword. + AsName, + + /// Occurs when we see a `.` between name-like tokens + /// after an `as` keyword. This implies we must parse + /// a `from` statement, since an `as` in a `from` can + /// never alias a dotted name. + AsDottedNameDot, + /// Occurs when we see a name-like token after a + /// `.name as`. + AsDottedName, + /// Occurs when we see a comma right before `a.b as foo`. + AsDottedNameComma, + /// Occurs before `, a.b as foo`. In this state, we can + /// see either a non-dotted alias or a dotted name. + AsDottedNameOrAlias, + /// Occurs before `bar, a.b as foo`. In this state, we can + /// see a `.`, `as`, or `import`. + AsDottedNameOrAliasName, + + /// Occurs when we've seen a dot right before the cursor + /// or after the first name-like token. That is, `.name`. + /// This could be from `import module.name` or `from ..name + /// import blah`. + InitialDot, + /// Occurs when we see `foo.bar`. When we enter + /// this state, it means we must be in an `import` + /// statement, since `from foo.bar` is always invalid. + InitialDotName, + /// Occurs when we see `.foo.bar`. This lets us + /// continue consuming a dotted name. + InitialDottedName, + + // When the states below occur, we are locked into + // recognizing a `from ... import ...` statement. + /// Occurs when we've seen an ellipsis right before the + /// cursor or after the first name-like token. That is, + /// `...name`. This must be from a + /// `from ...name import blah` statement. + FromEllipsisName, + /// A state for consuming `.` and `...` in a `from` import + /// statement. We enter this after seeing a `.` or a `...` + /// right after an `import` statement or a `...` right + /// before the end user's cursor. Either way, we have to + /// consume only dots at this point until we find a `from` + /// token. + FromDots, + /// Occurs when we've seen an `import` followed by a name-like + /// token. i.e., `from name import` or `from ...name import`. + FromDottedName, + /// Occurs when we've seen an `import` followed by a + /// name-like token with a dot. i.e., `from .name import` + /// or `from ..name import`. + FromDottedNameDot, + /// A `*` was just seen, which must mean the import is of + /// the form `from module import *`. + FromStar, + /// A left parenthesis was just seen. + FromLpar, + + // Below are terminal states. Once we reach one + // of these, the state machine ends. + /// We just saw a `from` token. We never have any + /// outgoing transitions from this. + From, + /// This is like `import`, but used in a context + /// where we know we're in an import statement and + /// specifically *not* a `from ... import ...` + /// statement. + ImportFinal, + } + + let mut state = S::Start; + // The token immediate before (or at) the cursor. + let last = tokens.last()?; + // A token corresponding to `import`, if found. + let mut import: Option<&Token> = None; + // A token corresponding to `from`, if found. + let mut from: Option<&Token> = None; + // Whether an initial dot was found right before the cursor, + // or right before the name at the cursor. + let mut initial_dot = false; + // An incomplete import statement was found. + // Usually either `from foo imp` + // or `import foo a`. + let mut incomplete_as_or_import = false; + for token in tokens.iter().rev().take(Self::LIMIT) { + if token.kind().is_trivia() { + continue; + } + state = match (state, token.kind()) { + // These cases handle our "initial" condition. + // Basically, this is what detects how to drop us into + // the full state machine below for parsing any kind of + // import statement. There are also some cases we try + // to detect here that indicate the user is probably + // typing an `import` or `as` token. In effect, we + // try to pluck off the initial name-like token that + // represents where the cursor likely is. And then + // we move on to try and detect the type of import + // statement that we're dealing with. + // (S::Start, TK::Newline) => S::Start, + (S::Start, TK::Star) => S::FromStar, + (S::Start, TK::Name) if typed.is_none() => S::AdjacentName, + (S::Start, TK::Name) => S::FirstName, + (S::Start | S::FirstName | S::AdjacentName, TK::Import) => S::Import, + (S::Start | S::FirstName | S::AdjacentName, TK::Lpar) => S::FromLpar, + (S::Start | S::FirstName | S::AdjacentName, TK::Comma) => S::NameList, + (S::Start | S::FirstName | S::AdjacentName, TK::Dot) => S::InitialDot, + (S::Start | S::FirstName | S::AdjacentName, TK::Ellipsis) => S::FromEllipsisName, + (S::Start | S::FirstName, TK::As) => S::As, + (S::Start | S::AdjacentName, TK::From) => S::From, + (S::FirstName, TK::From) => S::From, + (S::FirstName, TK::Name) => S::AdjacentName, + + // This handles the case where we see `.name`. Here, + // we could be in `from .name`, `from ..name`, `from + // ...name`, `from foo.name`, `import foo.name`, + // `import bar, foo.name` and so on. + (S::InitialDot, TK::Dot | TK::Ellipsis) => S::FromDots, + (S::InitialDot, TK::Name) => S::InitialDotName, + (S::InitialDot, TK::From) => S::From, + (S::InitialDotName, TK::Dot) => S::InitialDottedName, + (S::InitialDotName, TK::Ellipsis) => S::FromDots, + (S::InitialDotName, TK::As) => S::AsDottedNameDot, + (S::InitialDotName, TK::Comma) => S::AsDottedNameOrAlias, + (S::InitialDotName, TK::Import) => S::ImportFinal, + (S::InitialDotName, TK::From) => S::From, + (S::InitialDottedName, TK::Dot | TK::Ellipsis) => S::FromDots, + (S::InitialDottedName, TK::Name) => S::InitialDotName, + (S::InitialDottedName, TK::From) => S::From, + + // This state machine parses `dotted_as_names` or + // `import_from_as_names`. It has a carve out for when + // it finds a dot, which indicates it must parse only + // `dotted_as_names`. + (S::NameList, TK::Name | TK::Unknown) => S::NameListNameOrAlias, + (S::NameList, TK::Lpar) => S::FromLpar, + (S::NameListNameOrAlias, TK::As) => S::As, + (S::NameListNameOrAlias, TK::Comma) => S::NameList, + (S::NameListNameOrAlias, TK::Import) => S::Import, + (S::NameListNameOrAlias, TK::Lpar) => S::FromLpar, + (S::NameListNameOrAlias, TK::Unknown) => S::NameListNameOrAlias, + // This pops us out of generic name-list parsing + // and puts us firmly into `dotted_as_names` in + // the grammar. + (S::NameListNameOrAlias, TK::Dot) => S::AsDottedNameDot, + + // This identifies aliasing via `as`. The main trick + // here is that if we see a `.`, then we move to a + // different set of states since we know we must be in + // an `import` statement. Without a `.` though, we + // could be in an `import` or a `from`. For example, + // `import numpy as np` or + // `from collections import defaultdict as dd`. + (S::As, TK::Name) => S::AsName, + (S::AsName, TK::Dot) => S::AsDottedNameDot, + (S::AsName, TK::Import) => S::Import, + (S::AsName, TK::Comma) => S::NameList, + + // This is the mini state machine for handling + // `dotted_as_names`. We enter it when we see + // `foo.bar as baz`. We therefore know this must + // be an `import` statement and not a `from import` + // statement. + (S::AsDottedName, TK::Dot) => S::AsDottedNameDot, + (S::AsDottedName, TK::Comma) => S::AsDottedNameComma, + (S::AsDottedName, TK::Import) => S::ImportFinal, + (S::AsDottedNameDot, TK::Name) => S::AsDottedName, + (S::AsDottedNameComma, TK::Name) => S::AsDottedNameOrAlias, + (S::AsDottedNameOrAlias, TK::Name) => S::AsDottedNameOrAliasName, + (S::AsDottedNameOrAlias, TK::Dot) => S::AsDottedNameDot, + (S::AsDottedNameOrAliasName, TK::Dot | TK::As) => S::AsDottedNameDot, + (S::AsDottedNameOrAliasName, TK::Import) => S::ImportFinal, + + // A `*` and `(` immediately constrains what we're allowed to see. + // We can jump right to expecting an `import` keyword. + (S::FromStar | S::FromLpar, TK::Import) => S::Import, + + // The transitions below handle everything from `from` + // to `import`. Basically, once we see an `import` + // token or otherwise know we're parsing the module + // section of a `from` import statement, we end up in + // one of the transitions below. + (S::Import, TK::Dot | TK::Ellipsis) => S::FromDots, + (S::Import, TK::Name | TK::Unknown) => S::FromDottedName, + (S::FromDottedName, TK::Dot) => S::FromDottedNameDot, + (S::FromDottedName, TK::Ellipsis) => S::FromDots, + (S::FromDottedNameDot, TK::Name) => S::FromDottedName, + (S::FromDottedNameDot, TK::Dot | TK::Ellipsis) => S::FromDots, + (S::FromEllipsisName | S::FromDots, TK::Dot | TK::Ellipsis) => S::FromDots, + ( + S::FromEllipsisName | S::FromDots | S::FromDottedName | S::FromDottedNameDot, + TK::From, + ) => S::From, + + _ => break, + }; + // If we transition into a few different special + // states, we record the token. + match state { + S::Import | S::ImportFinal => { + import = Some(token); + } + S::From => { + from = Some(token); + } + S::AdjacentName => { + // We've seen two adjacent name-like tokens + // right before the cursor. At this point, + // we continue on to try to recognize a nearly + // valid import statement, and to figure out + // what kinds of completions we should offer + // (if any). + incomplete_as_or_import = true; + } + S::InitialDot | S::FromEllipsisName => { + initial_dot = true; + } + _ => {} + } + } + + // Now find a possibly dotted name up to where the current + // cursor is. This could be an item inside a module, a module + // name, a submodule name or even a relative module. The + // point is that it is the thing that the end user is trying + // to complete. + let source = source_text(db, file); + let mut to_complete = String::new(); + let end = last.range().end(); + let mut start = end; + for token in tokens.iter().rev().take(Self::LIMIT) { + match token.kind() { + TK::Name | TK::Dot | TK::Ellipsis => { + start = token.range().start(); + } + _ => break, + } + } + to_complete.push_str(&source[TextRange::new(start, end)]); + + // If the typed text corresponds precisely to a keyword, + // then as a special case, consider it "incomplete" for that + // keyword. This occurs when the cursor is immediately at the + // end of `import` or `as`, e.g., `import`. So we + // should provide it as a completion so that the end user can + // confirm it as-is. We special case this because a complete + // `import` or `as` gets special recognition as a special token + // kind, and it isn't worth complicating the state machine + // above to account for this. + // + // We also handle the more general "incomplete" cases here too. + // Basically, `incomplete_as_or_import` is set to `true` when + // we detect an "adjacent" name in an import statement. Some + // examples: + // + // from foo + // from foo imp + // from foo import bar + // from foo import bar a + // import foo + // import foo a + // + // Since there is a very limited number of cases, we can + // suggest `import` when an `import` token isn't present. And + // `as` when an `import` token *is* present. Notably, `as` can + // only appear after an `import` keyword! + if typed == Some("import") || (incomplete_as_or_import && import.is_none()) { + return Some(ImportStatement::Incomplete(IncompleteImport::Import)); + } else if typed == Some("as") || (incomplete_as_or_import && import.is_some()) { + return Some(ImportStatement::Incomplete(IncompleteImport::As)); + } + match (from, import) { + (None, None) => None, + (None, Some(import)) => { + let ast = find_ast_for_import(parsed, import)?; + // If we found a dot near the cursor, then this + // must be a request for submodule completions. + let kind = if initial_dot { + let (parent, _) = to_complete.rsplit_once('.')?; + let module_name = ModuleName::new(parent)?; + ImportKind::Submodule { + parent: module_name, + } + } else { + ImportKind::Module + }; + Some(ImportStatement::Import(Import { ast, kind })) + } + (Some(from), import) => { + let ast = find_ast_for_from_import(parsed, from)?; + let kind = if import.is_some() { + FromImportKind::Attribute + } else if initial_dot { + // TODO: Handle relative imports here. + if to_complete.starts_with('.') { + return Some(ImportStatement::Incomplete(IncompleteImport::Import)); + } + let (parent, _) = to_complete.rsplit_once('.')?; + let module_name = ModuleName::new(parent)?; + FromImportKind::Submodule { + parent: module_name, + } + } else { + FromImportKind::Module + }; + Some(ImportStatement::FromImport(FromImport { ast, kind })) + } + } + } + + /// Add completions, if any and if appropriate, based on the detected + /// import statement. + fn add_completions<'db>( + &self, + db: &'db dyn Db, + file: File, + completions: &mut Completions<'db>, + ) { + let model = SemanticModel::new(db, file); + match *self { + ImportStatement::Import(Import { ref kind, .. }) => match *kind { + ImportKind::Module => { + completions.extend(model.import_completions()); + } + ImportKind::Submodule { ref parent } => { + completions.extend(model.import_submodule_completions_for_name(parent)); + } + }, + ImportStatement::FromImport(FromImport { ast, ref kind }) => match *kind { + FromImportKind::Module => { + completions.extend(model.import_completions()); + } + FromImportKind::Submodule { ref parent } => { + completions.extend(model.import_submodule_completions_for_name(parent)); + } + FromImportKind::Attribute => { + completions.extend(model.from_import_completions(ast)); + } + }, + ImportStatement::Incomplete(IncompleteImport::As) => { + completions.try_add(Completion::keyword("as")); + } + ImportStatement::Incomplete(IncompleteImport::Import) => { + completions.try_add(Completion::keyword("import")); + } + } + } +} + +/// Finds the AST node, if available, corresponding to the given `from` +/// token. +/// +/// This always returns `None` when the `token` is not a `from` token. +fn find_ast_for_from_import<'p>( + parsed: &'p ParsedModuleRef, + token: &Token, +) -> Option<&'p ast::StmtImportFrom> { + let covering_node = covering_node(parsed.syntax().into(), token.range()) + .find_first(|node| node.is_stmt_import_from()) + .ok()?; + let ast::AnyNodeRef::StmtImportFrom(from_import) = covering_node.node() else { + return None; + }; + Some(from_import) +} + +/// Finds the AST node, if available, corresponding to the given `import` +/// token. +/// +/// This always returns `None` when the `token` is not a `import` token. +fn find_ast_for_import<'p>( + parsed: &'p ParsedModuleRef, + token: &Token, +) -> Option<&'p ast::StmtImport> { + let covering_node = covering_node(parsed.syntax().into(), token.range()) + .find_first(|node| node.is_stmt_import()) + .ok()?; + let ast::AnyNodeRef::StmtImport(import) = covering_node.node() else { + return None; + }; + Some(import) +} + /// Returns a slice of tokens that all start before the given /// [`TextSize`] offset. /// @@ -790,257 +1182,6 @@ fn token_suffix_by_kinds( })) } -/// Looks for the start of a `from module import ` 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 ` 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, - // - // - // 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 -} - -/// Looks for the start of a `import ` statement. -/// -/// This also handles cases like `import foo, c, bar`. -/// -/// If found, a token corresponding to the `import` or `from` keyword -/// and the closest point of the `` is returned. -/// -/// It is assumed that callers will call `from_import_tokens` first to -/// try and recognize a `from ... import ...` statement before using -/// this. -fn import_tokens(tokens: &[Token]) -> Option<(&Token, &Token)> { - use TokenKind as TK; - - /// A look-back limit, in order to bound work. - /// - /// See `LIMIT` in `import_from_tokens` for more context. - const LIMIT: usize = 1_000; - - /// A state used to "parse" the tokens preceding the user's cursor, - /// in reverse, to detect a `import` statement. - enum S { - Start, - Names, - } - - let mut state = S::Start; - let module_token = tokens.last()?; - // Move backward through the tokens until we get to - // the `import` 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. - (S::Start | S::Names, TK::Name | TK::Comma | TK::As | TK::Unknown) => S::Names, - (S::Start | S::Names, TK::Import | TK::From) => { - return Some((token, module_token)); - } - _ => return None, - }; - } - None -} - -/// Looks for the start of a `from module ` 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`. 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 -} - -/// 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. /// @@ -3790,7 +3931,7 @@ import } #[test] - fn import_multiple() { + fn import_multiple_betwixt() { let builder = completion_test_builder( "\ import re, c, sys @@ -3799,6 +3940,26 @@ import re, c, sys builder.build().contains("collections"); } + #[test] + fn import_multiple_end1() { + let builder = completion_test_builder( + "\ +import collections.abc, unico +", + ); + builder.build().contains("unicodedata"); + } + + #[test] + fn import_multiple_end2() { + let builder = completion_test_builder( + "\ +import collections.abc, urllib.parse, bu +", + ); + builder.build().contains("builtins"); + } + #[test] fn import_with_aliases() { let builder = completion_test_builder( @@ -4820,7 +4981,7 @@ from collections import defaultdict as f } #[test] - fn import_missing_alias_suggests_as() { + fn import_missing_alias_suggests_as_with_leading_char() { let builder = completion_test_builder( "\ import collections a @@ -4829,6 +4990,16 @@ import collections a assert_snapshot!(builder.build().snapshot(), @"as"); } + #[test] + fn import_missing_alias_suggests_as() { + let builder = completion_test_builder( + "\ +import collections + ", + ); + assert_snapshot!(builder.build().snapshot(), @"as"); + } + #[test] fn import_dotted_module_missing_alias_suggests_as() { let builder = completion_test_builder( @@ -4850,7 +5021,7 @@ import collections.abc as c, typing a } #[test] - fn from_import_missing_alias_suggests_as() { + fn from_import_missing_alias_suggests_as_with_leading_char() { let builder = completion_test_builder( "\ from collections.abc import Mapping a @@ -4859,6 +5030,16 @@ from collections.abc import Mapping a assert_snapshot!(builder.build().snapshot(), @"as"); } + #[test] + fn from_import_missing_alias_suggests_as() { + let builder = completion_test_builder( + "\ +from collections import defaultdict + ", + ); + assert_snapshot!(builder.build().snapshot(), @"as"); + } + #[test] fn from_import_parenthesized_missing_alias_suggests_as() { let builder = completion_test_builder( @@ -4991,7 +5172,7 @@ match status: } #[test] - fn from_import_import_suggests_nothing() { + fn from_import_import_suggests_import() { let builder = completion_test_builder("from typing import"); assert_snapshot!(builder.build().snapshot(), @"import"); } diff --git a/crates/ty_python_semantic/src/semantic_model.rs b/crates/ty_python_semantic/src/semantic_model.rs index 34efa22a02..d83f8dfe32 100644 --- a/crates/ty_python_semantic/src/semantic_model.rs +++ b/crates/ty_python_semantic/src/semantic_model.rs @@ -91,11 +91,7 @@ impl<'db> SemanticModel<'db> { } /// Returns completions for symbols available in a `from module import ` context. - pub fn from_import_completions( - &self, - import: &ast::StmtImportFrom, - _name: Option, - ) -> Vec> { + pub fn from_import_completions(&self, import: &ast::StmtImportFrom) -> Vec> { let module_name = match ModuleName::from_import_statement(self.db, self.file, import) { Ok(module_name) => module_name, Err(err) => { @@ -110,69 +106,8 @@ impl<'db> SemanticModel<'db> { self.module_completions(&module_name) } - /// Returns completions only for submodules for the module - /// identified by `name` in `import`. - /// - /// For example, `import re, os., zlib`. - pub fn import_submodule_completions( - &self, - import: &ast::StmtImport, - name: usize, - ) -> Vec> { - let module_ident = &import.names[name].name; - let Some((parent_ident, _)) = module_ident.rsplit_once('.') else { - return vec![]; - }; - let module_name = - match ModuleName::from_identifier_parts(self.db, self.file, Some(parent_ident), 0) { - Ok(module_name) => module_name, - Err(err) => { - tracing::debug!( - "Could not extract module name from `{module:?}`: {err:?}", - module = module_ident, - ); - return vec![]; - } - }; - self.import_submodule_completions_for_name(&module_name) - } - - /// Returns completions only for submodules for the module - /// used in a `from module import attribute` statement. - /// - /// For example, `from os.`. - pub fn from_import_submodule_completions( - &self, - import: &ast::StmtImportFrom, - ) -> Vec> { - let level = import.level; - let Some(module_ident) = import.module.as_deref() else { - return vec![]; - }; - let Some((parent_ident, _)) = module_ident.rsplit_once('.') else { - return vec![]; - }; - let module_name = match ModuleName::from_identifier_parts( - self.db, - self.file, - Some(parent_ident), - level, - ) { - Ok(module_name) => module_name, - Err(err) => { - tracing::debug!( - "Could not extract module name from `{module:?}` with level {level}: {err:?}", - module = import.module, - level = import.level, - ); - return vec![]; - } - }; - self.import_submodule_completions_for_name(&module_name) - } - /// Returns submodule-only completions for the given module. - fn import_submodule_completions_for_name( + pub fn import_submodule_completions_for_name( &self, module_name: &ModuleName, ) -> Vec> {