From 5f501374c4a0ef546b90131241f0850251acb8d8 Mon Sep 17 00:00:00 2001
From: Alex Waygood <Alex.Waygood@Gmail.com>
Date: Fri, 14 Nov 2025 13:11:49 +0000
Subject: [PATCH] [ty] Support attribute-expression `TYPE_CHECKING`
 conditionals (#21449)

---
 .../resources/mdtest/function/return_type.md  | 39 +++++++++++++++++++
 .../src/semantic_index/builder.rs             | 28 ++++++++++---
 2 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/crates/ty_python_semantic/resources/mdtest/function/return_type.md b/crates/ty_python_semantic/resources/mdtest/function/return_type.md
index 980afbdf46..284a2dd475 100644
--- a/crates/ty_python_semantic/resources/mdtest/function/return_type.md
+++ b/crates/ty_python_semantic/resources/mdtest/function/return_type.md
@@ -132,8 +132,29 @@ def f(x: int | str):
 Inside an `if TYPE_CHECKING` block, we allow "stub" style function definitions with empty bodies,
 since these functions will never actually be called.
 
+`compat/__init__.py`:
+
+```py
+```
+
+`compat/sub/__init__.py`:
+
+```py
+```
+
+`compat/sub/sub.py`:
+
 ```py
 from typing import TYPE_CHECKING
+```
+
+`main.py`:
+
+```py
+from typing import TYPE_CHECKING
+import typing
+import typing as t
+import compat.sub.sub
 
 if TYPE_CHECKING:
     def f() -> int: ...
@@ -199,6 +220,24 @@ if get_bool():
 if TYPE_CHECKING:
     if not TYPE_CHECKING:
         def n() -> str: ...
+
+if typing.TYPE_CHECKING:
+    def o() -> str: ...
+
+if not typing.TYPE_CHECKING:
+    def p() -> str: ...  # error: [invalid-return-type]
+
+if compat.sub.sub.TYPE_CHECKING:
+    def q() -> str: ...
+
+if not compat.sub.sub.TYPE_CHECKING:
+    def r() -> str: ...  # error: [invalid-return-type]
+
+if t.TYPE_CHECKING:
+    def s() -> str: ...
+
+if not t.TYPE_CHECKING:
+    def t() -> str: ...  # error: [invalid-return-type]
 ```
 
 ## Conditional return type
diff --git a/crates/ty_python_semantic/src/semantic_index/builder.rs b/crates/ty_python_semantic/src/semantic_index/builder.rs
index cc1b1649fa..17dd2c6c2e 100644
--- a/crates/ty_python_semantic/src/semantic_index/builder.rs
+++ b/crates/ty_python_semantic/src/semantic_index/builder.rs
@@ -3131,15 +3131,31 @@ impl ExpressionsScopeMapBuilder {
 
 /// 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")
+    fn is_dotted_name(expr: &ast::Expr) -> bool {
+        match expr {
+            ast::Expr::Name(_) => true,
+            ast::Expr::Attribute(ast::ExprAttribute { value, .. }) => is_dotted_name(value),
+            _ => false,
+        }
+    }
+
+    match expr {
+        ast::Expr::Name(ast::ExprName { id, .. }) => id == "TYPE_CHECKING",
+        ast::Expr::Attribute(ast::ExprAttribute { value, attr, .. }) => {
+            attr == "TYPE_CHECKING" && is_dotted_name(value)
+        }
+        _ => false,
+    }
 }
 
 /// 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"
-        )
+    matches!(
+        expr,
+        ast::Expr::UnaryOp(ast::ExprUnaryOp {
+            op: ast::UnaryOp::Not,
+            operand,
+            ..
+        }) if is_if_type_checking(operand)
     )
 }