[ty] More fully support 'if typing.TYPE_CHECKING'

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

@@ -134,6 +134,7 @@ since these functions will never actually be called.
 
 ```py
 from typing import TYPE_CHECKING
+import typing
 
 if TYPE_CHECKING:
     def f() -> int: ...
@@ -199,6 +200,9 @@ if get_bool():
 if TYPE_CHECKING:
     if not TYPE_CHECKING:
         def n() -> str: ...
+
+if typing.TYPE_CHECKING:
+    def o() -> str: ...
 ```
 
 ## Conditional return type

crates/ty_python_semantic/resources/mdtest/known_constants.md

@@ -3,27 +3,37 @@
 ## `typing.TYPE_CHECKING`
 
 This constant is `True` when in type-checking mode, `False` otherwise. The symbol is defined to be
-`False` at runtime. In typeshed, it is annotated as `bool`. This test makes sure that we infer
-`Literal[True]` for it anyways.
+`False` at runtime. In typeshed, it is annotated as `bool`.
 
 ### Basic
 
 ```py
 from typing import TYPE_CHECKING
-import typing
 
-reveal_type(TYPE_CHECKING)  # revealed: Literal[True]
-reveal_type(typing.TYPE_CHECKING)  # revealed: Literal[True]
-```
+if TYPE_CHECKING:
+    type_checking = True
+if not TYPE_CHECKING:
+    runtime = True
 
-### Aliased
+# type_checking is treated as unconditionally assigned.
+reveal_type(type_checking)  # revealed: Literal[True]
+# error: [unresolved-reference]
+reveal_type(runtime)  # revealed: Unknown
+```
 
-Make sure that we still infer the correct type if the constant has been given a different name:
+### As module attribute
 
 ```py
-from typing import TYPE_CHECKING as TC
+import typing
 
-reveal_type(TC)  # revealed: Literal[True]
+if typing.TYPE_CHECKING:
+    type_checking = True
+if not typing.TYPE_CHECKING:
+    runtime = True
+
+reveal_type(type_checking)  # revealed: Literal[True]
+# error: [unresolved-reference]
+reveal_type(runtime)  # revealed: Unknown
 ```
 
 ### `typing_extensions` re-export
@@ -33,7 +43,14 @@ This should behave in the same way as `typing.TYPE_CHECKING`:
 ```py
 from typing_extensions import TYPE_CHECKING
 
-reveal_type(TYPE_CHECKING)  # revealed: Literal[True]
+if TYPE_CHECKING:
+    type_checking = True
+if not TYPE_CHECKING:
+    runtime = True
+
+reveal_type(type_checking)  # revealed: Literal[True]
+# error: [unresolved-reference]
+reveal_type(runtime)  # revealed: Unknown
 ```
 
 ## User-defined `TYPE_CHECKING`
@@ -46,7 +63,7 @@ type checkers, e.g. mypy and pyright.
 
 ```py
 TYPE_CHECKING = False
-reveal_type(TYPE_CHECKING)  # revealed: Literal[True]
+
 if TYPE_CHECKING:
     type_checking = True
 if not TYPE_CHECKING:
@@ -61,11 +78,11 @@ reveal_type(runtime)  # revealed: Unknown
 ### With a type annotation
 
 We can also define `TYPE_CHECKING` with a type annotation. The type must be one to which `bool` can
-be assigned. Even in this case, the type of `TYPE_CHECKING` is still inferred to be `Literal[True]`.
+be assigned.
 
 ```py
 TYPE_CHECKING: bool = False
-reveal_type(TYPE_CHECKING)  # revealed: Literal[True]
+
 if TYPE_CHECKING:
     type_checking = True
 if not TYPE_CHECKING:
@@ -84,6 +101,21 @@ reveal_type(runtime)  # revealed: Unknown
 TYPE_CHECKING = False
 ```
 
+```py
+from constants import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    type_checking = True
+if not TYPE_CHECKING:
+    runtime = True
+
+reveal_type(type_checking)  # revealed: Literal[True]
+# error: [unresolved-reference]
+reveal_type(runtime)  # revealed: Unknown
+```
+
+### Importing user-defined `TYPE_CHECKING` from stub
+
 `stub.pyi`:
 
 ```pyi
@@ -93,13 +125,16 @@ TYPE_CHECKING: bool = ...
 ```
 
 ```py
-from constants import TYPE_CHECKING
-
-reveal_type(TYPE_CHECKING)  # revealed: Literal[True]
-
 from stub import TYPE_CHECKING
 
-reveal_type(TYPE_CHECKING)  # revealed: Literal[True]
+if TYPE_CHECKING:
+    type_checking = True
+if not TYPE_CHECKING:
+    runtime = True
+
+reveal_type(type_checking)  # revealed: Literal[True]
+# error: [unresolved-reference]
+reveal_type(runtime)  # revealed: Unknown
 ```
 
 ### Invalid assignment to `TYPE_CHECKING`
@@ -122,12 +157,14 @@ TYPE_CHECKING: int = 1
 # error: [invalid-type-checking-constant]
 TYPE_CHECKING: str = "str"
 
+# error: [invalid-assignment]
 # error: [invalid-type-checking-constant]
 TYPE_CHECKING: str = False
 
 # error: [invalid-type-checking-constant]
 TYPE_CHECKING: Literal[False] = False
 
+# error: [invalid-assignment]
 # error: [invalid-type-checking-constant]
 TYPE_CHECKING: Literal[True] = False
 ```
@@ -140,6 +177,7 @@ from typing import Literal
 # error: [invalid-type-checking-constant]
 TYPE_CHECKING: str
 
+# error: [invalid-assignment]
 # error: [invalid-type-checking-constant]
 TYPE_CHECKING: str = False
 

crates/ty_python_semantic/resources/mdtest/type_api.md

@@ -154,7 +154,6 @@ the expression is not of statically known truthiness.
 
 ```py
 from ty_extensions import static_assert
-from typing import TYPE_CHECKING
 import sys
 
 static_assert(True)
@@ -174,8 +173,6 @@ static_assert("d" in "abc")  # error: "Static assertion error: argument evaluate
 n = None
 static_assert(n is None)
 
-static_assert(TYPE_CHECKING)
-
 static_assert(sys.version_info >= (3, 6))
 ```
 

crates/ty_python_semantic/src/place.rs

@@ -754,14 +754,10 @@ fn place_by_id<'db>(
             // a diagnostic if we see it being modified externally. In type inference, we
             // can assign a "narrow" type to it even if it is not *declared*. This means, we
             // do not have to call [`widen_type_for_undeclared_public_symbol`].
-            //
-            // `TYPE_CHECKING` is a special variable that should only be assigned `False`
-            // at runtime, but is always considered `True` in type checking.
-            // See mdtest/known_constants.md#user-defined-type_checking for details.
             let is_considered_non_modifiable = place_table(db, scope)
                 .place_expr(place_id)
                 .expr
-                .is_name_and(|name| matches!(name, "__slots__" | "TYPE_CHECKING"));
+                .is_name_and(|name| matches!(name, "__slots__"));
 
             if scope.file(db).is_stub(db) {
                 // We generally trust module-level undeclared places in stubs and do not union

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -549,14 +549,20 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
     }
 
     fn build_predicate(&mut self, predicate_node: &ast::Expr) -> PredicateOrLiteral<'db> {
-        // Some commonly used test expressions are eagerly evaluated as `true`
-        // or `false` here for performance reasons. This list does not need to
-        // be exhaustive. More complex expressions will still evaluate to the
-        // correct value during type-checking.
+        // Some commonly used test expressions are eagerly evaluated as `true` or `false` here for
+        // performance reasons. This list does not need to be exhaustive. More complex expressions
+        // will still evaluate to the correct value during type-checking. (The one exception is
+        // `TYPE_CHECKING`; we need to detect it here in order to handle it correctly in
+        // conditions; in type inference it will resolve to its runtime value.)
         fn resolve_to_literal(node: &ast::Expr) -> Option<bool> {
             match node {
                 ast::Expr::BooleanLiteral(ast::ExprBooleanLiteral { value, .. }) => Some(*value),
                 ast::Expr::Name(ast::ExprName { id, .. }) if id == "TYPE_CHECKING" => Some(true),
+                ast::Expr::Attribute(ast::ExprAttribute { attr, .. })
+                    if attr == "TYPE_CHECKING" =>
+                {
+                    Some(true)
+                }
                 ast::Expr::NumberLiteral(ast::ExprNumberLiteral {
                     value: ast::Number::Int(n),
                     ..
@@ -2753,14 +2759,12 @@ 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")
+        || matches!(expr, ast::Expr::Attribute(ast::ExprAttribute { attr, .. }) if attr == "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"
-        )
+        && is_if_type_checking(operand)
     )
 }

crates/ty_python_semantic/src/types/infer.rs

@@ -3900,7 +3900,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                     ) {
                         report_invalid_type_checking_constant(&self.context, target.into());
                     }
-                    Type::BooleanLiteral(true)
+                    value_ty
                 } else if self.in_stub() && value.is_ellipsis_literal_expr() {
                     Type::unknown()
                 } else {
@@ -3989,7 +3989,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 // otherwise, assigning something other than `False` is an error
                 report_invalid_type_checking_constant(&self.context, target.into());
             }
-            declared_ty.inner = Type::BooleanLiteral(true);
         }
 
         // Handle various singletons.
@@ -4014,12 +4013,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
         if let Some(value) = value {
             let inferred_ty = self.infer_expression(value);
-            let inferred_ty = if target
-                .as_name_expr()
-                .is_some_and(|name| &name.id == "TYPE_CHECKING")
-            {
-                Type::BooleanLiteral(true)
-            } else if self.in_stub() && value.is_ellipsis_literal_expr() {
+            let inferred_ty = if self.in_stub() && value.is_ellipsis_literal_expr() {
                 declared_ty.inner_type()
             } else {
                 inferred_ty