[ty] Consider type_check_only when ranking completions (#20910)

Author: decorator-factory

crates/ty_completion_eval/completion-evaluation-tasks.csv

@@ -4,6 +4,11 @@ higher-level-symbols-preferred,main.py,0,
 higher-level-symbols-preferred,main.py,1,1
 import-deprioritizes-dunder,main.py,0,1
 import-deprioritizes-sunder,main.py,0,1
+import-deprioritizes-type_check_only,main.py,0,1
+import-deprioritizes-type_check_only,main.py,1,1
+import-deprioritizes-type_check_only,main.py,2,1
+import-deprioritizes-type_check_only,main.py,3,2
+import-deprioritizes-type_check_only,main.py,4,3
 internal-typeshed-hidden,main.py,0,4
 none-completion,main.py,0,11
 numpy-array,main.py,0,

crates/ty_completion_eval/truth/import-deprioritizes-type_check_only/completion.toml

@@ -0,0 +1,2 @@
+[settings]
+auto-import = true

crates/ty_completion_eval/truth/import-deprioritizes-type_check_only/main.py

@@ -0,0 +1,12 @@
+from module import UniquePrefixA<CURSOR:UniquePrefixAzurous>
+from module import unique_prefix_<CURSOR:unique_prefix_azurous>
+
+from module import Class
+
+Class.meth_<CURSOR:meth_azurous>
+
+# TODO: bound methods don't preserve type-check-only-ness, this is a bug
+Class().meth_<CURSOR:meth_azurous>
+
+# TODO: auto-imports don't take type-check-only-ness into account, this is a bug
+UniquePrefixA<CURSOR:module.UniquePrefixAzurous>

crates/ty_completion_eval/truth/import-deprioritizes-type_check_only/module.py

@@ -0,0 +1,20 @@
+from typing import type_check_only
+
+
+@type_check_only
+class UniquePrefixApple: pass
+
+class UniquePrefixAzurous: pass
+
+
+@type_check_only
+def unique_prefix_apple() -> None: pass
+
+def unique_prefix_azurous() -> None: pass
+
+
+class Class:
+    @type_check_only
+    def meth_apple(self) -> None: pass
+
+    def meth_azurous(self) -> None: pass

crates/ty_completion_eval/truth/import-deprioritizes-type_check_only/pyproject.toml

@@ -0,0 +1,5 @@
+[project]
+name = "test"
+version = "0.1.0"
+requires-python = ">=3.13"
+dependencies = []

crates/ty_completion_eval/truth/import-deprioritizes-type_check_only/uv.lock

@@ -65,6 +65,9 @@ pub struct Completion<'db> {
     /// use it mainly in tests so that we can write less
     /// noisy tests.
     pub builtin: bool,
+    /// Whether this item only exists for type checking purposes and
+    /// will be missing at runtime
+    pub is_type_check_only: bool,
     /// The documentation associated with this item, if
     /// available.
     pub documentation: Option<Docstring>,
@@ -79,6 +82,7 @@ impl<'db> Completion<'db> {
             .ty
             .and_then(|ty| DefinitionsOrTargets::from_ty(db, ty));
         let documentation = definition.and_then(|def| def.docstring(db));
+        let is_type_check_only = semantic.is_type_check_only(db);
         Completion {
             name: semantic.name,
             insert: None,
@@ -87,6 +91,7 @@ impl<'db> Completion<'db> {
             module_name: None,
             import: None,
             builtin: semantic.builtin,
+            is_type_check_only,
             documentation,
         }
     }
@@ -294,6 +299,7 @@ fn add_keyword_value_completions<'db>(
             kind: None,
             module_name: None,
             import: None,
+            is_type_check_only: false,
             builtin: true,
             documentation: None,
         });
@@ -339,6 +345,8 @@ fn add_unimported_completions<'db>(
             module_name: Some(symbol.module.name(db)),
             import: import_action.import().cloned(),
             builtin: false,
+            // TODO: `is_type_check_only` requires inferring the type of the symbol
+            is_type_check_only: false,
             documentation: None,
         });
     }
@@ -837,16 +845,21 @@ fn is_in_string(parsed: &ParsedModuleRef, offset: TextSize) -> bool {
     })
 }
 
-/// Order completions lexicographically, with these exceptions:
+/// Order completions according to the following rules:
 ///
-/// 1) A `_[^_]` prefix sorts last and
-/// 2) A `__` prefix sorts last except before (1)
+/// 1) Names with no underscore prefix
+/// 2) Names starting with `_` but not dunders
+/// 3) `__dunder__` names
+///
+/// Among each category, type-check-only items are sorted last,
+/// and otherwise completions are sorted lexicographically.
 ///
 /// This has the effect of putting all dunder attributes after "normal"
 /// attributes, and all single-underscore attributes after dunder attributes.
 fn compare_suggestions(c1: &Completion, c2: &Completion) -> Ordering {
     let (kind1, kind2) = (NameKind::classify(&c1.name), NameKind::classify(&c2.name));
-    kind1.cmp(&kind2).then_with(|| c1.name.cmp(&c2.name))
+
+    (kind1, c1.is_type_check_only, &c1.name).cmp(&(kind2, c2.is_type_check_only, &c2.name))
 }
 
 #[cfg(test)]
@@ -3398,6 +3411,65 @@ from os.<CURSOR>
         );
     }
 
+    #[test]
+    fn import_type_check_only_lowers_ranking() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                r#"
+                import foo
+                foo.A<CURSOR>
+                "#,
+            )
+            .source(
+                "foo/__init__.py",
+                r#"
+                from typing import type_check_only
+
+                @type_check_only
+                class Apple: pass
+
+                class Banana: pass
+                class Cat: pass
+                class Azorubine: pass
+                "#,
+            )
+            .build();
+
+        let settings = CompletionSettings::default();
+        let completions = completion(&test.db, &settings, test.cursor.file, test.cursor.offset);
+
+        let [apple_pos, banana_pos, cat_pos, azo_pos, ann_pos] =
+            ["Apple", "Banana", "Cat", "Azorubine", "__annotations__"].map(|name| {
+                completions
+                    .iter()
+                    .position(|comp| comp.name == name)
+                    .unwrap()
+            });
+
+        assert!(completions[apple_pos].is_type_check_only);
+        assert!(apple_pos > banana_pos.max(cat_pos).max(azo_pos));
+        assert!(ann_pos > apple_pos);
+    }
+
+    #[test]
+    fn type_check_only_is_type_check_only() {
+        // `@typing.type_check_only` is a function that's unavailable at runtime
+        // and so should be the last "non-underscore" completion in `typing`
+        let test = cursor_test("from typing import t<CURSOR>");
+
+        let settings = CompletionSettings::default();
+        let completions = completion(&test.db, &settings, test.cursor.file, test.cursor.offset);
+        let last_nonunderscore = completions
+            .into_iter()
+            .filter(|c| !c.name.starts_with('_'))
+            .next_back()
+            .unwrap();
+
+        assert_eq!(&last_nonunderscore.name, "type_check_only");
+        assert!(last_nonunderscore.is_type_check_only);
+    }
+
     #[test]
     fn regression_test_issue_642() {
         // Regression test for https://github.com/astral-sh/ty/issues/642

crates/ty_ide/src/completion.rs

@@ -342,6 +342,12 @@ pub struct Completion<'db> {
     pub builtin: bool,
 }
 
+impl<'db> Completion<'db> {
+    pub fn is_type_check_only(&self, db: &'db dyn Db) -> bool {
+        self.ty.is_some_and(|ty| ty.is_type_check_only(db))
+    }
+}
+
 pub trait HasType {
     /// Returns the inferred type of `self`.
     ///

crates/ty_python_semantic/src/semantic_model.rs

@@ -53,8 +53,8 @@ pub use crate::types::display::DisplaySettings;
 use crate::types::display::TupleSpecialization;
 use crate::types::enums::{enum_metadata, is_single_member_enum};
 use crate::types::function::{
-    DataclassTransformerFlags, DataclassTransformerParams, FunctionSpans, FunctionType,
-    KnownFunction,
+    DataclassTransformerFlags, DataclassTransformerParams, FunctionDecorators, FunctionSpans,
+    FunctionType, KnownFunction,
 };
 pub(crate) use crate::types::generics::GenericContext;
 use crate::types::generics::{
@@ -868,6 +868,17 @@ impl<'db> Type<'db> {
         matches!(self, Type::Dynamic(_))
     }
 
+    /// Is a value of this type only usable in typing contexts?
+    pub(crate) fn is_type_check_only(&self, db: &'db dyn Db) -> bool {
+        match self {
+            Type::ClassLiteral(class_literal) => class_literal.type_check_only(db),
+            Type::FunctionLiteral(f) => {
+                f.has_known_decorator(db, FunctionDecorators::TYPE_CHECK_ONLY)
+            }
+            _ => false,
+        }
+    }
+
     // If the type is a specialized instance of the given `KnownClass`, returns the specialization.
     pub(crate) fn known_specialization(
         &self,

crates/ty_python_semantic/src/types.rs

@@ -1001,6 +1001,7 @@ impl<'db> Bindings<'db> {
                                     class_literal.body_scope(db),
                                     class_literal.known(db),
                                     class_literal.deprecated(db),
+                                    class_literal.type_check_only(db),
                                     Some(params),
                                     class_literal.dataclass_transformer_params(db),
                                 )));