[ty] Preserve typing-only completion ranking in TYPE_CHECKING blocks

AGENTS.md

@@ -111,6 +111,18 @@ When the task matches a more specific ty workflow, also read and follow that ski
 - Ecosystem report summaries: `.agents/skills/summarise-ecosystem-results/SKILL.md`.
 - Reproducing, investigating, or minimizing ecosystem or primer differences: `.agents/skills/minimizing-ty-ecosystem-changes/SKILL.md`.
 
+### Completion ranking
+
+When changing ty autocomplete ranking, add or update evaluation fixtures under `crates/ty_completion_eval/truth/`. Extend an existing project when it is a good fit for the behavior being tested; otherwise, add a new one. Use `<CURSOR:expected_name>` directives to assert ranking, and include the expected module for auto-import completions. Add `completion.rs` unit tests only when the evaluation fixtures cannot adequately cover the behavior.
+
+Regenerate and review the committed evaluation results after changing ranking behavior or fixtures:
+
+```sh
+CARGO_PROFILE_DEV_OPT_LEVEL=1 CARGO_PROFILE_DEV_LTO=off CARGO_PROFILE_DEV_DEBUG="line-tables-only" cargo run --package ty_completion_eval -- all --threshold 0.4 --tasks crates/ty_completion_eval/completion-evaluation-tasks.csv
+```
+
+To inspect one evaluation task, run `cargo run --package ty_completion_eval -- show-one <fixture-name> --file-name <file-name> --index <cursor-index>`.
+
 ### Ad hoc reproductions
 
 When running ty against a temporary Python reproduction file, create it outside the Ruff checkout (for example, under `/tmp`). A file inside the checkout discovers Ruff's root `pyproject.toml`, whose `requires-python = ">=3.7"` causes ty to infer Python 3.7 as the default Python version.

crates/ty_completion_eval/completion-evaluation-tasks.csv

@@ -19,8 +19,20 @@ 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
+import-deprioritizes-type_check_only,main.py,3,1
+import-deprioritizes-type_check_only,main.py,4,1
+import-deprioritizes-type_check_only,main.py,5,2
+import-deprioritizes-type_check_only,main.py,6,3
+import-deprioritizes-type_check_only,main.py,7,1
+import-deprioritizes-type_check_only,main.py,8,1
+import-deprioritizes-type_check_only,main.py,9,1
+import-deprioritizes-type_check_only,main.py,10,1
+import-deprioritizes-type_check_only,main.py,11,1
+import-deprioritizes-type_check_only,main.py,12,1
+import-deprioritizes-type_check_only,main.py,13,1
+import-deprioritizes-type_check_only,main.py,14,1
+import-deprioritizes-type_check_only,main.py,15,1
+import-deprioritizes-type_check_only,main.py,16,1
 import-deprioritizes-type_check_only,main.pyi,0,1
 import-deprioritizes-type_check_only,main.pyi,1,1
 import-keyword-completion,main.py,0,1
@@ -54,6 +66,9 @@ typing-only-auto-import-ranking,main.py,2,1
 typing-only-auto-import-ranking,main.py,3,2
 typing-only-auto-import-ranking,main.py,4,1
 typing-only-auto-import-ranking,main.py,5,1
+typing-only-auto-import-ranking,main.py,6,1
+typing-only-auto-import-ranking,main.py,7,1
+typing-only-auto-import-ranking,main.py,8,1
 typing-only-auto-import-ranking,main.pyi,0,1
 typing-only-auto-import-ranking,main.pyi,1,1
 typing-only-auto-import-ranking,main.pyi,2,1

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

@@ -1,12 +1,48 @@
+from typing import TYPE_CHECKING
+
+import private_stub
+
 from module import UniquePrefixA<CURSOR:UniquePrefixAzurous>
 from module import unique_prefix_<CURSOR:unique_prefix_azurous>
+from private_stub import _Al<CURSOR:_Alzeta>
 
 from module import Class
 
 Class.meth_<CURSOR:meth_azurous>
+private_stub._Al<CURSOR:_Alzeta>
 
 # 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>
+
+if TYPE_CHECKING:
+    from module import UniquePrefixA<CURSOR:UniquePrefixApple>
+    from module import unique_prefix_<CURSOR:unique_prefix_apple>
+    from private_stub import _Al<CURSOR:_Alpha>
+
+    Class.meth_<CURSOR:meth_apple>
+    private_stub._Al<CURSOR:_Alpha>
+
+    def declared_in_type_checking_block() -> None:
+        private_stub._Al<CURSOR:_Alpha>
+
+
+def function_scope() -> None:
+    if TYPE_CHECKING:
+        from private_stub import _Al<CURSOR:_Alpha>
+
+        private_stub._Al<CURSOR:_Alpha>
+
+
+if not TYPE_CHECKING:
+    pass
+else:
+    private_stub._Al<CURSOR:_Alpha>
+
+
+if TYPE_CHECKING:
+    pass
+else:
+    from private_stub import _Al<CURSOR:_Alzeta>

crates/ty_completion_eval/truth/import-deprioritizes-type_check_only/private_stub.pyi

@@ -0,0 +1,4 @@
+from typing import TypeVar
+
+_Alpha = TypeVar("_Alpha")
+_Alzeta = 1

crates/ty_completion_eval/truth/typing-only-auto-import-ranking/main.py

@@ -1,3 +1,5 @@
+from typing import TYPE_CHECKING
+
 # Runtime symbols outrank alternatives from typing-only modules in Python files.
 deprecated<CURSOR: warnings.deprecated>
 NoneTy<CURSOR: types.NoneType>
@@ -7,3 +9,13 @@
 static_ass<CURSOR: ty_extensions.static_assert>
 is_equiv<CURSOR: ty_extensions._internal.is_equivalent_to>
 TypedDictFall<CURSOR: _typeshed._type_checker_internals.TypedDictFallback>
+
+# Typing-only symbols retain their usual ranking inside TYPE_CHECKING blocks.
+if TYPE_CHECKING:
+    deprecated<CURSOR: typing_extensions.deprecated>
+    NoneTy<CURSOR: _typeshed.NoneType>
+
+
+def function_scope() -> None:
+    if TYPE_CHECKING:
+        deprecated<CURSOR: typing_extensions.deprecated>

crates/ty_ide/src/completion.rs

@@ -1,3 +1,4 @@
+use std::cell::OnceCell;
 use std::cmp::Ordering;
 use std::collections::{BinaryHeap, binary_heap};
 use ty_python_semantic::ProgramEnvironment;
@@ -10,15 +11,15 @@ use ruff_python_ast::find_node::{CoveringNode, covering_node};
 use ruff_python_ast::name::{Name, UnqualifiedName};
 use ruff_python_ast::str::Quote;
 use ruff_python_ast::token::{Token, TokenKind, Tokens};
-use ruff_python_ast::{self as ast, AnyNodeRef, PySourceType};
+use ruff_python_ast::{self as ast, AnyNodeRef};
 use ruff_python_codegen::Stylist;
 use ruff_python_literal::escape::{Escape, UnicodeEscape};
 use ruff_text_size::{Ranged, TextRange, TextSize};
 use rustc_hash::FxHashSet;
 use ty_module_resolver::{
     ImportingFile, KnownModule, Module, ModuleName, resolve_real_shadowable_module,
 };
-use ty_python_core::ProgramFile;
+use ty_python_core::{ProgramFile, semantic_index};
 use ty_python_semantic::HasType;
 use ty_python_semantic::types::{SpecialFormType, UnionType};
 use ty_python_semantic::{
@@ -525,12 +526,7 @@ impl<'db> CompletionBuilder<'db> {
         let kind = self
             .kind
             .or_else(|| self.ty.and_then(|ty| completion_kind_from_type(db, ty)));
-        let relevance = Relevance::new(
-            collection_context,
-            query,
-            &self,
-            program_file.file(db).source_type(db),
-        );
+        let relevance = Relevance::new(db, program_file, collection_context, query, &self);
         let (label, insert, insert_text_format, command) =
             if collection_context.should_complete_callable_parentheses(kind) {
                 let label = self.insert.unwrap_or_else(|| self.name.clone());
@@ -838,8 +834,13 @@ impl<'m> Context<'m> {
         settings: &CompletionSettings,
         capabilities: CompletionCapabilities,
     ) -> CollectionContext<'db> {
+        let type_checking_block = Some(TypeCheckingBlock::at_cursor(self.cursor.range));
+
         match self.kind {
-            ContextKind::Keywords(_) | ContextKind::Import(_) => CollectionContext::none(),
+            ContextKind::Keywords(_) | ContextKind::Import(_) => CollectionContext {
+                type_checking_block,
+                ..CollectionContext::none()
+            },
             ContextKind::NonImport(_) => {
                 let env = model.program_environment();
                 let exception_ty = self.cursor.exception_ty(db, &env);
@@ -859,6 +860,7 @@ impl<'m> Context<'m> {
                 CollectionContext {
                     exception_ty,
                     is_raising_exception: exception_ty.is_some(),
+                    type_checking_block,
                     complete_class_parentheses: complete_callable_parentheses
                         && existing_class_bases.is_none()
                         && !self.cursor.suppress_class_parentheses(model),
@@ -1630,6 +1632,40 @@ impl UserQuery {
     }
 }
 
+#[derive(Clone, Debug)]
+struct TypeCheckingBlock {
+    range: TextRange,
+    is_inside: OnceCell<bool>,
+}
+
+impl TypeCheckingBlock {
+    fn at_cursor(range: TextRange) -> Self {
+        Self {
+            range,
+            is_inside: OnceCell::new(),
+        }
+    }
+
+    fn is_inside<'db>(&self, db: &'db dyn Db, file: ProgramFile<'db>) -> bool {
+        // Most completions are ranked independently of `TYPE_CHECKING`, so only query the
+        // semantic index when a typing-only completion needs to know the cursor's context.
+        *self.is_inside.get_or_init(|| {
+            let parsed = parsed_module(db, file.python_file(db)).load(db);
+            let index = semantic_index(db, file);
+
+            covering_node(parsed.syntax().into(), self.range)
+                .ancestors()
+                .find_map(|node| {
+                    let ast::AnyNodeRef::StmtIf(statement) = node else {
+                        return None;
+                    };
+                    index.try_expression_scope_id(statement.test.as_ref())
+                })
+                .is_some_and(|scope| index.is_in_type_checking_block(scope, self.range))
+        })
+    }
+}
+
 /// Context used to help filter completions when collecting them.
 #[derive(Clone, Debug, Default)]
 struct CollectionContext<'db> {
@@ -1640,6 +1676,8 @@ struct CollectionContext<'db> {
     exception_ty: Option<Type<'db>>,
     /// Whether we're in an exception context (`raise` or `except`) or not.
     is_raising_exception: bool,
+    /// Whether the cursor is inside a type-checking-only block, if a cursor is available.
+    type_checking_block: Option<TypeCheckingBlock>,
     /// Names of base classes that are already specified in the class definition,
     /// including the class being defined (unless its name was previously bound).
     /// Used to filter out duplicate and self-referential base class suggestions.
@@ -1795,11 +1833,12 @@ impl Relevance {
     ///
     /// A smaller rank means the completion should appear higher in the
     /// results shown to end users.
-    fn new(
-        _ctx: &CollectionContext,
+    fn new<'db>(
+        db: &'db dyn Db,
+        program_file: ProgramFile<'db>,
+        ctx: &CollectionContext,
         query: &UserQuery,
         c: &CompletionBuilder,
-        source_type: PySourceType,
     ) -> Relevance {
         Relevance {
             definitively_usable: if c.is_context_specific {
@@ -1837,7 +1876,13 @@ impl Relevance {
             } else {
                 Sort::Even
             },
-            type_check_only: if c.is_type_check_only && !source_type.is_stub() {
+            type_check_only: if c.is_type_check_only
+                && !program_file.file(db).source_type(db).is_stub()
+                && !ctx
+                    .type_checking_block
+                    .as_ref()
+                    .is_some_and(|block| block.is_inside(db, program_file))
+            {
                 Sort::Lower
             } else {
                 Sort::Even