From 6bc33a041f71963c1bcb83150d0432e54ae1631e Mon Sep 17 00:00:00 2001
From: Andrew Gallant <andrew@astral.sh>
Date: Tue, 2 Sep 2025 09:40:56 -0400
Subject: [PATCH] [ty] Pick up typed text as a query for unimported completions

It's almost certainly bad juju to show literally every single possible
symbol when completions are requested but there is nothing typed yet.
Moreover, since there are so many symbols, it is likely beneficial to
try and winnow them down before sending them to the client.

This change tries to extract text that has been typed and then uses
that as a query to listing all available symbols.
---
 crates/ty_ide/src/completion.rs | 67 ++++++++++++++++++++++++++-------
 1 file changed, 54 insertions(+), 13 deletions(-)

diff --git a/crates/ty_ide/src/completion.rs b/crates/ty_ide/src/completion.rs
index 80239b18e1..8ef7e02e32 100644
--- a/crates/ty_ide/src/completion.rs
+++ b/crates/ty_ide/src/completion.rs
@@ -47,16 +47,18 @@ pub fn completion<'db>(
         CompletionTargetAst::Import { .. } | CompletionTargetAst::ImportViaFrom { .. } => {
             model.import_completions()
         }
-        CompletionTargetAst::Scoped { node } => {
+        CompletionTargetAst::Scoped { node, typed } => {
             let mut completions = model.scoped_completions(node);
             if settings.auto_import {
-                for symbol in all_symbols(db, "") {
-                    completions.push(Completion {
-                        name: ast::name::Name::new(&symbol.symbol.name),
-                        ty: None,
-                        kind: symbol.symbol.kind.to_completion_kind(),
-                        builtin: false,
-                    });
+                if let Some(typed) = typed {
+                    for symbol in all_symbols(db, typed) {
+                        completions.push(Completion {
+                            name: ast::name::Name::new(&symbol.symbol.name),
+                            ty: None,
+                            kind: symbol.symbol.kind.to_completion_kind(),
+                            builtin: false,
+                        });
+                    }
                 }
             }
             completions
@@ -79,10 +81,12 @@ pub fn completion<'db>(
         .collect()
 }
 
+#[derive(Clone, Debug)]
 pub struct DetailedCompletion<'db> {
     pub inner: Completion<'db>,
     pub documentation: Option<Docstring>,
 }
+
 impl<'db> std::ops::Deref for DetailedCompletion<'db> {
     type Target = Completion<'db>;
     fn deref(&self) -> &Self::Target {
@@ -283,16 +287,22 @@ impl<'t> CompletionTargetTokens<'t> {
                 }
             }
             CompletionTargetTokens::Generic { token } => {
-                let covering_node = covering_node(parsed.syntax().into(), token.range());
-                Some(CompletionTargetAst::Scoped {
-                    node: covering_node.node(),
-                })
+                let node = covering_node(parsed.syntax().into(), token.range()).node();
+                let typed = match node {
+                    ast::AnyNodeRef::ExprName(ast::ExprName { id, .. }) => {
+                        let name = id.as_str();
+                        if name.is_empty() { None } else { Some(name) }
+                    }
+                    _ => None,
+                };
+                Some(CompletionTargetAst::Scoped { node, typed })
             }
             CompletionTargetTokens::Unknown => {
                 let range = TextRange::empty(offset);
                 let covering_node = covering_node(parsed.syntax().into(), range);
                 Some(CompletionTargetAst::Scoped {
                     node: covering_node.node(),
+                    typed: None,
                 })
             }
         }
@@ -346,7 +356,16 @@ enum CompletionTargetAst<'t> {
     },
     /// 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> },
+    Scoped {
+        /// The node with the smallest range that fully covers
+        /// the token under the cursor.
+        node: ast::AnyNodeRef<'t>,
+        /// The text that has been typed so far, if available.
+        ///
+        /// When not `None`, the typed text is guaranteed to be
+        /// non-empty.
+        typed: Option<&'t str>,
+    },
 }
 
 /// Returns a suffix of `tokens` corresponding to the `kinds` given.
@@ -3038,6 +3057,24 @@ from os.<CURSOR>
         test.assert_completions_do_not_include("abspath");
     }
 
+    #[test]
+    fn auto_import_with_submodule() {
+        let test = CursorTest::builder()
+            .source("main.py", "Abra<CURSOR>")
+            .source("package/__init__.py", "AbraKadabra = 1")
+            .build();
+
+        let settings = CompletionSettings { auto_import: true };
+        let expected = "AbraKadabra";
+        let completions = completion(&test.db, &settings, test.cursor.file, test.cursor.offset);
+        assert!(
+            completions
+                .iter()
+                .any(|completion| completion.name == expected),
+            "Expected completions to include `{expected}`"
+        );
+    }
+
     #[test]
     fn regression_test_issue_642() {
         // Regression test for https://github.com/astral-sh/ty/issues/642
@@ -3056,6 +3093,10 @@ from os.<CURSOR>
         );
     }
 
+    // NOTE: The methods below are getting somewhat ridiculous.
+    // We should refactor this by converting to using a builder
+    // to set different modes. ---AG
+
     impl CursorTest {
         /// Returns all completions except for builtins.
         fn completions_without_builtins(&self) -> String {