[red-knot] Add rule invalid-type-checking-constant (#16501)

## Summary

This PR adds more features to #16468.

* Adds a new error rule `invalid-type-checking-constant`, which occurs
when we try to assign a value other than `False` to a user-defined
`TYPE_CHECKING` variable (it is possible to assign `...` in a stub
file).
* Allows annotated assignment to `TYPE_CHECKING`. Only types that
`False` can be assigned to are allowed. However, the type of
`TYPE_CHECKING` will be inferred to be `Literal[True]` regardless of
what the type is specified.

## 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>

crates/red_knot_python_semantic/resources/mdtest/known_constants.md

@@ -26,12 +26,24 @@ from typing import TYPE_CHECKING as TC
 reveal_type(TC)  # revealed: Literal[True]
 ```
 
-### User-defined `TYPE_CHECKING`
+### `typing_extensions` re-export
+
+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]
+```
+
+## User-defined `TYPE_CHECKING`
 
 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.
 
+### With no type annotation
+
 ```py
 TYPE_CHECKING = False
 reveal_type(TYPE_CHECKING)  # revealed: Literal[True]
@@ -46,6 +58,24 @@ reveal_type(type_checking)  # revealed: Literal[True]
 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]`.
+
+```py
+TYPE_CHECKING: bool = False
+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
+```
+
 ### Importing user-defined `TYPE_CHECKING`
 
 `constants.py`:
@@ -54,19 +84,68 @@ reveal_type(runtime)  # revealed: Unknown
 TYPE_CHECKING = False
 ```
 
+`stub.pyi`:
+
+```pyi
+TYPE_CHECKING: bool
+# or
+TYPE_CHECKING: bool = ...
+```
+
 ```py
 from constants import TYPE_CHECKING
 
-# constants.TYPE_CHECKING is modifiable, but it is still treated as True.
+reveal_type(TYPE_CHECKING)  # revealed: Literal[True]
+
+from stub import TYPE_CHECKING
+
 reveal_type(TYPE_CHECKING)  # revealed: Literal[True]
 ```
 
-### `typing_extensions` re-export
+### Invalid assignment to `TYPE_CHECKING`
 
-This should behave in the same way as `typing.TYPE_CHECKING`:
+Only `False` can be assigned to `TYPE_CHECKING`; any assignment other than `False` will result in an
+error. A type annotation to which `bool` is not assignable is also an error.
 
 ```py
-from typing_extensions import TYPE_CHECKING
+from typing import Literal
 
-reveal_type(TYPE_CHECKING)  # revealed: Literal[True]
+# error: [invalid-type-checking-constant]
+TYPE_CHECKING = True
+
+# error: [invalid-type-checking-constant]
+TYPE_CHECKING: bool = True
+
+# error: [invalid-type-checking-constant]
+TYPE_CHECKING: int = 1
+
+# error: [invalid-type-checking-constant]
+TYPE_CHECKING: str = "str"
+
+# error: [invalid-type-checking-constant]
+TYPE_CHECKING: str = False
+
+# error: [invalid-type-checking-constant]
+TYPE_CHECKING: Literal[False] = False
+
+# error: [invalid-type-checking-constant]
+TYPE_CHECKING: Literal[True] = False
+```
+
+The same rules apply in a stub file:
+
+```pyi
+from typing import Literal
+
+# error: [invalid-type-checking-constant]
+TYPE_CHECKING: str
+
+# error: [invalid-type-checking-constant]
+TYPE_CHECKING: str = False
+
+# error: [invalid-type-checking-constant]
+TYPE_CHECKING: Literal[False] = ...
+
+# error: [invalid-type-checking-constant]
+TYPE_CHECKING: object = "str"
 ```

crates/red_knot_python_semantic/src/symbol.rs

@@ -504,14 +504,6 @@ fn symbol_impl<'db>(
 ) -> Symbol<'db> {
     let _span = tracing::trace_span!("symbol", ?name).entered();
 
-    // We don't need to check for `typing_extensions` here, because `typing_extensions.TYPE_CHECKING`
-    // is just a re-export of `typing.TYPE_CHECKING`.
-    if name == "TYPE_CHECKING"
-        && file_to_module(db, scope.file(db))
-            .is_some_and(|module| module.is_known(KnownModule::Typing))
-    {
-        return Symbol::bound(Type::BooleanLiteral(true));
-    }
     if name == "platform"
         && file_to_module(db, scope.file(db))
             .is_some_and(|module| module.is_known(KnownModule::Sys))

crates/red_knot_python_semantic/src/types/diagnostic.rs

@@ -39,6 +39,7 @@ pub(crate) fn register_lints(registry: &mut LintRegistryBuilder) {
     registry.register_lint(&INVALID_METACLASS);
     registry.register_lint(&INVALID_PARAMETER_DEFAULT);
     registry.register_lint(&INVALID_RAISE);
+    registry.register_lint(&INVALID_TYPE_CHECKING_CONSTANT);
     registry.register_lint(&INVALID_TYPE_FORM);
     registry.register_lint(&INVALID_TYPE_VARIABLE_CONSTRAINTS);
     registry.register_lint(&MISSING_ARGUMENT);
@@ -412,6 +413,24 @@ declare_lint! {
     }
 }
 
+declare_lint! {
+    /// ## What it does
+    /// Checks for a value other than `False` assigned to the `TYPE_CHECKING` variable, or an
+    /// annotation not assignable from `bool`.
+    ///
+    /// ## Why is this bad?
+    /// The name `TYPE_CHECKING` is reserved for a flag that can be used to provide conditional
+    /// code seen only by the type checker, and not at runtime. Normally this flag is imported from
+    /// `typing` or `typing_extensions`, but it can also be defined locally. If defined locally, it
+    /// must be assigned the value `False` at runtime; the type checker will consider its value to
+    /// be `True`. If annotated, it must be annotated as a type that can accept `bool` values.
+    pub(crate) static INVALID_TYPE_CHECKING_CONSTANT = {
+        summary: "detects invalid TYPE_CHECKING constant assignments",
+        status: LintStatus::preview("1.0.0"),
+        default_level: Level::Error,
+    }
+}
+
 declare_lint! {
     /// ## What it does
     /// Checks for invalid type expressions.
@@ -1042,6 +1061,14 @@ pub(super) fn report_invalid_attribute_assignment(
     );
 }
 
+pub(super) fn report_invalid_type_checking_constant(context: &InferContext, node: AnyNodeRef) {
+    context.report_lint(
+        &INVALID_TYPE_CHECKING_CONSTANT,
+        node,
+        format_args!("The name TYPE_CHECKING is reserved for use as a flag; only False can be assigned to it.",),
+    );
+}
+
 pub(super) fn report_possibly_unresolved_reference(
     context: &InferContext,
     expr_name_node: &ast::ExprName,

crates/red_knot_python_semantic/src/types/infer.rs

@@ -83,9 +83,10 @@ use super::class_base::ClassBase;
 use super::context::{InNoTypeCheck, InferContext, WithDiagnostics};
 use super::diagnostic::{
     report_index_out_of_bounds, report_invalid_exception_caught, report_invalid_exception_cause,
-    report_invalid_exception_raised, report_non_subscriptable,
-    report_possibly_unresolved_reference, report_slice_step_size_zero, report_unresolved_reference,
-    INVALID_METACLASS, STATIC_ASSERT_ERROR, SUBCLASS_OF_FINAL_CLASS, TYPE_ASSERTION_FAILURE,
+    report_invalid_exception_raised, report_invalid_type_checking_constant,
+    report_non_subscriptable, report_possibly_unresolved_reference, report_slice_step_size_zero,
+    report_unresolved_reference, INVALID_METACLASS, STATIC_ASSERT_ERROR, SUBCLASS_OF_FINAL_CLASS,
+    TYPE_ASSERTION_FAILURE,
 };
 use super::slots::check_class_slots;
 use super::string_annotation::{
@@ -2153,6 +2154,15 @@ impl<'db> TypeInferenceBuilder<'db> {
                 // 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" {
+                    if !matches!(
+                        value.as_boolean_literal_expr(),
+                        Some(ast::ExprBooleanLiteral { value: false, .. })
+                    ) {
+                        report_invalid_type_checking_constant(
+                            &self.context,
+                            assignment.name().into(),
+                        );
+                    }
                     Type::BooleanLiteral(true)
                 } else if self.in_stub() && value.is_ellipsis_literal_expr() {
                     Type::unknown()
@@ -2209,6 +2219,34 @@ impl<'db> TypeInferenceBuilder<'db> {
             DeferredExpressionState::from(self.are_all_types_deferred()),
         );
 
+        if target
+            .as_name_expr()
+            .is_some_and(|name| &name.id == "TYPE_CHECKING")
+        {
+            if !KnownClass::Bool
+                .to_instance(self.db())
+                .is_assignable_to(self.db(), declared_ty.inner_type())
+            {
+                // annotation not assignable from `bool` is an error
+                report_invalid_type_checking_constant(&self.context, assignment.into());
+            } else if self.in_stub()
+                && value
+                    .as_ref()
+                    .is_none_or(|value| value.is_ellipsis_literal_expr())
+            {
+                // stub file assigning nothing or `...` is fine
+            } else if !matches!(
+                value
+                    .as_ref()
+                    .and_then(|value| value.as_boolean_literal_expr()),
+                Some(ast::ExprBooleanLiteral { value: false, .. })
+            ) {
+                // otherwise, assigning something other than `False` is an error
+                report_invalid_type_checking_constant(&self.context, assignment.into());
+            }
+            declared_ty.inner = Type::BooleanLiteral(true);
+        }
+
         // Handle various singletons.
         if let Type::Instance(instance) = declared_ty.inner_type() {
             if instance
@@ -2229,7 +2267,12 @@ impl<'db> TypeInferenceBuilder<'db> {
 
         if let Some(value) = value.as_deref() {
             let inferred_ty = self.infer_expression(value);
-            let inferred_ty = if self.in_stub() && value.is_ellipsis_literal_expr() {
+            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() {
                 declared_ty.inner_type()
             } else {
                 inferred_ty

knot.schema.json

@@ -451,6 +451,16 @@
             }
           ]
         },
+        "invalid-type-checking-constant": {
+          "title": "detects invalid TYPE_CHECKING constant assignments",
+          "description": "## What it does\nChecks for a value other than `False` assigned to the `TYPE_CHECKING` variable, or an\nannotation not assignable from `bool`.\n\n## Why is this bad?\nThe name `TYPE_CHECKING` is reserved for a flag that can be used to provide conditional\ncode seen only by the type checker, and not at runtime. Normally this flag is imported from\n`typing` or `typing_extensions`, but it can also be defined locally. If defined locally, it\nmust be assigned the value `False` at runtime; the type checker will consider its value to\nbe `True`. If annotated, it must be annotated as a type that can accept `bool` values.",
+          "default": "error",
+          "oneOf": [
+            {
+              "$ref": "#/definitions/Level"
+            }
+          ]
+        },
         "invalid-type-form": {
           "title": "detects invalid type forms",
           "description": "## What it does\nChecks for invalid type expressions.\n\n## Why is this bad?\nTODO #14889",