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

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

@@ -0,0 +1,8 @@
+version = 1
+revision = 3
+requires-python = ">=3.13"
+
+[[package]]
+name = "test"
+version = "0.1.0"
+source = { virtual = "." }

crates/ty_ide/src/completion.rs

@@ -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
+/// 1) Names with no underscore prefix
-/// 2) A `__` prefix sorts last except before (1)
+/// 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_python_semantic/src/semantic_model.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/types.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,
+    DataclassTransformerFlags, DataclassTransformerParams, FunctionDecorators, FunctionSpans,
-    KnownFunction,
+    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/call/bind.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),
                                 )));

crates/ty_python_semantic/src/types/class.rs

@@ -1336,6 +1336,8 @@ pub struct ClassLiteral<'db> {
     /// If this class is deprecated, this holds the deprecation message.
     pub(crate) deprecated: Option<DeprecatedInstance<'db>>,
 
+    pub(crate) type_check_only: bool,
+
     pub(crate) dataclass_params: Option<DataclassParams<'db>>,
     pub(crate) dataclass_transformer_params: Option<DataclassTransformerParams<'db>>,
 }

crates/ty_python_semantic/src/types/function.rs

@@ -121,6 +121,8 @@ bitflags! {
         const STATICMETHOD = 1 << 5;
         /// `@typing.override`
         const OVERRIDE = 1 << 6;
+        /// `@typing.type_check_only`
+        const TYPE_CHECK_ONLY = 1 << 7;
     }
 }
 
@@ -135,6 +137,7 @@ impl FunctionDecorators {
                 Some(KnownFunction::AbstractMethod) => FunctionDecorators::ABSTRACT_METHOD,
                 Some(KnownFunction::Final) => FunctionDecorators::FINAL,
                 Some(KnownFunction::Override) => FunctionDecorators::OVERRIDE,
+                Some(KnownFunction::TypeCheckOnly) => FunctionDecorators::TYPE_CHECK_ONLY,
                 _ => FunctionDecorators::empty(),
             },
             Type::ClassLiteral(class) => match class.known(db) {
@@ -1256,6 +1259,8 @@ pub enum KnownFunction {
     DisjointBase,
     /// [`typing(_extensions).no_type_check`](https://typing.python.org/en/latest/spec/directives.html#no-type-check)
     NoTypeCheck,
+    /// `typing(_extensions).type_check_only`
+    TypeCheckOnly,
 
     /// `typing(_extensions).assert_type`
     AssertType,
@@ -1340,7 +1345,7 @@ impl KnownFunction {
             .then_some(candidate)
     }
 
-    /// Return `true` if `self` is defined in `module` at runtime.
+    /// Return `true` if `self` is defined in `module`
     const fn check_module(self, module: KnownModule) -> bool {
         match self {
             Self::IsInstance
@@ -1394,6 +1399,8 @@ impl KnownFunction {
             | Self::NegatedRangeConstraint
             | Self::AllMembers => module.is_ty_extensions(),
             Self::ImportModule => module.is_importlib(),
+
+            Self::TypeCheckOnly => matches!(module, KnownModule::Typing),
         }
     }
 
@@ -1799,6 +1806,8 @@ pub(crate) mod tests {
                 | KnownFunction::DisjointBase
                 | KnownFunction::NoTypeCheck => KnownModule::TypingExtensions,
 
+                KnownFunction::TypeCheckOnly => KnownModule::Typing,
+
                 KnownFunction::IsSingleton
                 | KnownFunction::IsSubtypeOf
                 | KnownFunction::GenericContext

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -2207,6 +2207,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         let known_function =
             KnownFunction::try_from_definition_and_name(self.db(), definition, name);
 
+        // `type_check_only` is itself not available at runtime
+        if known_function == Some(KnownFunction::TypeCheckOnly) {
+            function_decorators |= FunctionDecorators::TYPE_CHECK_ONLY;
+        }
+
         let body_scope = self
             .index
             .node_scope(NodeWithScopeRef::Function(function))
@@ -2649,6 +2654,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         } = class_node;
 
         let mut deprecated = None;
+        let mut type_check_only = false;
         let mut dataclass_params = None;
         let mut dataclass_transformer_params = None;
         for decorator in decorator_list {
@@ -2673,6 +2679,14 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 continue;
             }
 
+            if decorator_ty
+                .as_function_literal()
+                .is_some_and(|function| function.is_known(self.db(), KnownFunction::TypeCheckOnly))
+            {
+                type_check_only = true;
+                continue;
+            }
+
             if let Type::FunctionLiteral(f) = decorator_ty {
                 // We do not yet detect or flag `@dataclass_transform` applied to more than one
                 // overload, or an overload and the implementation both. Nevertheless, this is not
@@ -2721,6 +2735,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 body_scope,
                 maybe_known_class,
                 deprecated,
+                type_check_only,
                 dataclass_params,
                 dataclass_transformer_params,
             )),