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

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,
-    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/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,
             )),