[ty] Support empty function bodies in if TYPE_CHECKING blocks (#19372)

## Summary

Resolves astral-sh/ty#339

Supports having a blank function body inside `if TYPE_CHECKING` block or
in the elif or else of a `if not TYPE_CHECKING` block.

```py
if TYPE_CHECKING:
    def foo() -> int: ...

if not TYPE_CHECKING: ...
else:     
    def bar() -> int: ...
```

## Test Plan

Update `function/return_type.md`

---------

Co-authored-by: Carl Meyer <[email protected]>

Author: MatthewMckee4

crates/ty_python_semantic/resources/mdtest/function/return_type.md

@@ -127,6 +127,80 @@ def f(x: int | str):
     return x
 ```
 
+### In `if TYPE_CHECKING` block
+
+Inside an `if TYPE_CHECKING` block, we allow "stub" style function definitions with empty bodies,
+since these functions will never actually be called.
+
+```py
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    def f() -> int: ...
+
+else:
+    def f() -> str:
+        return "hello"
+
+reveal_type(f)  # revealed: def f() -> int
+
+if not TYPE_CHECKING:
+    ...
+elif True:
+    def g() -> str: ...
+
+else:
+    def h() -> str: ...
+
+if not TYPE_CHECKING:
+    def i() -> int:
+        return 1
+
+else:
+    def i() -> str: ...
+
+reveal_type(i)  # revealed: def i() -> str
+
+if False:
+    ...
+elif TYPE_CHECKING:
+    def j() -> str: ...
+
+else:
+    def j_() -> str: ...  # error: [invalid-return-type]
+
+if False:
+    ...
+elif not TYPE_CHECKING:
+    def k_() -> str: ...  # error: [invalid-return-type]
+
+else:
+    def k() -> str: ...
+
+class Foo:
+    if TYPE_CHECKING:
+        def f(self) -> int: ...
+
+if TYPE_CHECKING:
+    class Bar:
+        def f(self) -> int: ...
+
+def get_bool() -> bool:
+    return True
+
+if TYPE_CHECKING:
+    if get_bool():
+        def l() -> str: ...
+
+if get_bool():
+    if TYPE_CHECKING:
+        def m() -> str: ...
+
+if TYPE_CHECKING:
+    if not TYPE_CHECKING:
+        def n() -> str: ...
+```
+
 ## Conditional return type
 
 ```py

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -90,6 +90,8 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> {
 
     /// Flags about the file's global scope
     has_future_annotations: bool,
+    /// Whether we are currently visiting an `if TYPE_CHECKING` block.
+    in_type_checking_block: bool,
 
     // Used for checking semantic syntax errors
     python_version: PythonVersion,
@@ -130,6 +132,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             try_node_context_stack_manager: TryNodeContextStackManager::default(),
 
             has_future_annotations: false,
+            in_type_checking_block: false,
 
             scopes: IndexVec::new(),
             place_tables: IndexVec::new(),
@@ -248,6 +251,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             node_with_kind,
             children_start..children_start,
             reachability,
+            self.in_type_checking_block,
         );
         let is_class_scope = scope.kind().is_class();
         self.try_node_context_stack_manager.enter_nested_scope();
@@ -719,7 +723,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         // since its the pattern that introduces any constraints, not the body.) Ideally, that
         // standalone expression would wrap the match arm's pattern as a whole. But a standalone
         // expression can currently only wrap an ast::Expr, which patterns are not. So, we need to
-        // choose an Expr that can “stand in” for the pattern, which we can wrap in a standalone
+        // choose an Expr that can "stand in" for the pattern, which we can wrap in a standalone
         // expression.
         //
         // See the comment in TypeInferenceBuilder::infer_match_pattern for more details.
@@ -1498,6 +1502,17 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                 let mut last_predicate = self.record_expression_narrowing_constraint(&node.test);
                 let mut last_reachability_constraint =
                     self.record_reachability_constraint(last_predicate);
+
+                let is_outer_block_in_type_checking = self.in_type_checking_block;
+
+                let if_block_in_type_checking = is_if_type_checking(&node.test);
+
+                // Track if we're in a chain that started with "not TYPE_CHECKING"
+                let mut is_in_not_type_checking_chain = is_if_not_type_checking(&node.test);
+
+                self.in_type_checking_block =
+                    if_block_in_type_checking || is_outer_block_in_type_checking;
+
                 self.visit_body(&node.body);
 
                 let mut post_clauses: Vec<FlowSnapshot> = vec![];
@@ -1516,6 +1531,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                     // if there's no `else` branch, we should add a no-op `else` branch
                     Some((None, Default::default()))
                 });
+
                 for (clause_test, clause_body) in elif_else_clauses {
                     // snapshot after every block except the last; the last one will just become
                     // the state that we merge the other snapshots into
@@ -1538,12 +1554,34 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                             self.record_reachability_constraint(last_predicate);
                     }
 
+                    // Determine if this clause is in type checking context
+                    let clause_in_type_checking = if let Some(elif_test) = clause_test {
+                        if is_if_type_checking(elif_test) {
+                            // This block has "TYPE_CHECKING" condition
+                            true
+                        } else if is_if_not_type_checking(elif_test) {
+                            // This block has "not TYPE_CHECKING" condition so we update the chain state for future blocks
+                            is_in_not_type_checking_chain = true;
+                            false
+                        } else {
+                            // This block has some other condition
+                            // It's in type checking only if we're in a "not TYPE_CHECKING" chain
+                            is_in_not_type_checking_chain
+                        }
+                    } else {
+                        is_in_not_type_checking_chain
+                    };
+
+                    self.in_type_checking_block = clause_in_type_checking;
+
                     self.visit_body(clause_body);
                 }
 
                 for post_clause_state in post_clauses {
                     self.flow_merge(post_clause_state);
                 }
+
+                self.in_type_checking_block = is_outer_block_in_type_checking;
             }
             ast::Stmt::While(ast::StmtWhile {
                 test,
@@ -2711,3 +2749,18 @@ impl ExpressionsScopeMapBuilder {
         ExpressionsScopeMap(interval_map.into_boxed_slice())
     }
 }
+
+/// Returns if the expression is a `TYPE_CHECKING` expression.
+fn is_if_type_checking(expr: &ast::Expr) -> bool {
+    matches!(expr, ast::Expr::Name(ast::ExprName { id, .. }) if id == "TYPE_CHECKING")
+}
+
+/// Returns if the expression is a `not TYPE_CHECKING` expression.
+fn is_if_not_type_checking(expr: &ast::Expr) -> bool {
+    matches!(expr, ast::Expr::UnaryOp(ast::ExprUnaryOp { op, operand, .. }) if *op == ruff_python_ast::UnaryOp::Not
+        && matches!(
+            &**operand,
+            ast::Expr::Name(ast::ExprName { id, .. }) if id == "TYPE_CHECKING"
+        )
+    )
+}

crates/ty_python_semantic/src/semantic_index/place.rs

@@ -525,10 +525,20 @@ impl FileScopeId {
 
 #[derive(Debug, salsa::Update, get_size2::GetSize)]
 pub struct Scope {
+    /// The parent scope, if any.
     parent: Option<FileScopeId>,
+
+    /// The node that introduces this scope.
     node: NodeWithScopeKind,
+
+    /// The range of [`FileScopeId`]s that are descendants of this scope.
     descendants: Range<FileScopeId>,
+
+    /// The constraint that determines the reachability of this scope.
     reachability: ScopedReachabilityConstraintId,
+
+    /// Whether this scope is defined inside an `if TYPE_CHECKING:` block.
+    in_type_checking_block: bool,
 }
 
 impl Scope {
@@ -537,12 +547,14 @@ impl Scope {
         node: NodeWithScopeKind,
         descendants: Range<FileScopeId>,
         reachability: ScopedReachabilityConstraintId,
+        in_type_checking_block: bool,
     ) -> Self {
         Scope {
             parent,
             node,
             descendants,
             reachability,
+            in_type_checking_block,
         }
     }
 
@@ -573,6 +585,10 @@ impl Scope {
     pub(crate) fn reachability(&self) -> ScopedReachabilityConstraintId {
         self.reachability
     }
+
+    pub(crate) fn in_type_checking_block(&self) -> bool {
+        self.in_type_checking_block
+    }
 }
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]

crates/ty_python_semantic/src/types/infer.rs

@@ -2135,6 +2135,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 if self.in_function_overload_or_abstractmethod() {
                     return;
                 }
+                if self.scope().scope(self.db()).in_type_checking_block() {
+                    return;
+                }
                 if let Some(class) = self.class_context_of_current_method() {
                     enclosing_class_context = Some(class);
                     if class.is_protocol(self.db()) {