From 1977dda079ffc2bb7eea4e4c904c2b87624b2645 Mon Sep 17 00:00:00 2001
From: Shunsuke Shibayama <45118249+mtshiba@users.noreply.github.com>
Date: Wed, 5 Mar 2025 00:58:29 +0900
Subject: [PATCH] [red-knot] respect `TYPE_CHECKING` even if not imported from
 typing (#16468)

## Summary

This PR closes #15722.

The change is that if the variable `TYPE_CHECKING` is defined/imported,
the type of the variable is interpreted as `Literal[True]` regardless of
what the value is.
This is compatible with the behavior of other type checkers (e.g. mypy,
pyright).

## Test Plan

I ran the tests with `cargo test -p red_knot_python_semantic` and
confirmed that all tests passed.

---------

Co-authored-by: Carl Meyer <carl@astral.sh>
---
 .../resources/mdtest/known_constants.md       | 28 +++++++++++++++----
 crates/red_knot_python_semantic/src/symbol.rs |  8 ++++--
 .../src/types/infer.rs                        |  7 ++++-
 3 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/crates/red_knot_python_semantic/resources/mdtest/known_constants.md b/crates/red_knot_python_semantic/resources/mdtest/known_constants.md
index b25de94d96..971f8f925f 100644
--- a/crates/red_knot_python_semantic/resources/mdtest/known_constants.md
+++ b/crates/red_knot_python_semantic/resources/mdtest/known_constants.md
@@ -26,21 +26,39 @@ from typing import TYPE_CHECKING as TC
 reveal_type(TC)  # revealed: Literal[True]
 ```
 
-### Must originate from `typing`
+### User-defined `TYPE_CHECKING`
 
-Make sure we only use our special handling for `typing.TYPE_CHECKING` and not for other constants
-with the same name:
+If we set `TYPE_CHECKING = False` directly instead of importing it from the `typing` module, it will
+still be treated as `True` during type checking. This behavior is for compatibility with other major
+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:
+    runtime = True
+
+# type_checking is treated as unconditionally assigned.
+reveal_type(type_checking)  # revealed: Literal[True]
+# error: [unresolved-reference]
+reveal_type(runtime)  # revealed: Unknown
+```
+
+### Importing user-defined `TYPE_CHECKING`
 
 `constants.py`:
 
 ```py
-TYPE_CHECKING: bool = False
+TYPE_CHECKING = False
 ```
 
 ```py
 from constants import TYPE_CHECKING
 
-reveal_type(TYPE_CHECKING)  # revealed: bool
+# constants.TYPE_CHECKING is modifiable, but it is still treated as True.
+reveal_type(TYPE_CHECKING)  # revealed: Literal[True]
 ```
 
 ### `typing_extensions` re-export
diff --git a/crates/red_knot_python_semantic/src/symbol.rs b/crates/red_knot_python_semantic/src/symbol.rs
index 1513f130ec..68598d65a5 100644
--- a/crates/red_knot_python_semantic/src/symbol.rs
+++ b/crates/red_knot_python_semantic/src/symbol.rs
@@ -458,8 +458,12 @@ fn symbol_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`].
-            let is_considered_non_modifiable =
-                symbol_table(db, scope).symbol(symbol_id).name() == "__slots__";
+            // `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 = symbol_table(db, scope).symbol(symbol_id).name()
+                == "__slots__"
+                || symbol_table(db, scope).symbol(symbol_id).name() == "TYPE_CHECKING";
 
             widen_type_for_undeclared_public_symbol(db, inferred, is_considered_non_modifiable)
                 .into()
diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
index 224086ec44..792916f836 100644
--- a/crates/red_knot_python_semantic/src/types/infer.rs
+++ b/crates/red_knot_python_semantic/src/types/infer.rs
@@ -2149,7 +2149,12 @@ impl<'db> TypeInferenceBuilder<'db> {
                 unpacked.expression_type(name_ast_id)
             }
             TargetKind::Name => {
-                if self.in_stub() && value.is_ellipsis_literal_expr() {
+                // `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.
+                if &name.id == "TYPE_CHECKING" {
+                    Type::BooleanLiteral(true)
+                } else if self.in_stub() && value.is_ellipsis_literal_expr() {
                     Type::unknown()
                 } else {
                     value_ty