Recognize all symbols named TYPE_CHECKING for in_type_checking_block (#15719)

Closes #15681

## Summary

This changes `analyze::typing::is_type_checking_block` to recognize all
symbols named "TYPE_CHECKING".
This matches the current behavior of mypy and pyright as well as
`flake8-type-checking`.

It also drops support for detecting `if False:` and `if 0:` as type
checking blocks. This used to be an option for
providing backwards compatibility with Python versions that did not have
a `typing` module, but has since
been removed from the typing spec and is no longer supported by any of
the mainstream type checkers.

## Test Plan

`cargo nextest run`

---------

Co-authored-by: Micha Reiser <[email protected]>

Author: Daverball

crates/ruff_linter/src/checkers/ast/mod.rs

@@ -248,6 +248,11 @@ impl<'a> Checker<'a> {
         cell_offsets: Option<&'a CellOffsets>,
         notebook_index: Option<&'a NotebookIndex>,
     ) -> Checker<'a> {
+        let mut semantic = SemanticModel::new(&settings.typing_modules, path, module);
+        if settings.preview.is_enabled() {
+            // Set the feature flag to test `TYPE_CHECKING` semantic changes
+            semantic.flags |= SemanticModelFlags::NEW_TYPE_CHECKING_BLOCK_DETECTION;
+        }
         Self {
             parsed,
             parsed_type_annotation: None,
@@ -263,7 +268,7 @@ impl<'a> Checker<'a> {
             stylist,
             indexer,
             importer: Importer::new(parsed, locator, stylist),
-            semantic: SemanticModel::new(&settings.typing_modules, path, module),
+            semantic,
             visit: deferred::Visit::default(),
             analyze: deferred::Analyze::default(),
             diagnostics: Vec::default(),

crates/ruff_linter/src/importer/mod.rs

@@ -9,7 +9,7 @@ use anyhow::Result;
 use libcst_native::{ImportAlias, Name as cstName, NameOrAttribute};
 
 use ruff_diagnostics::Edit;
-use ruff_python_ast::{self as ast, ModModule, Stmt};
+use ruff_python_ast::{self as ast, Expr, ModModule, Stmt};
 use ruff_python_codegen::Stylist;
 use ruff_python_parser::{Parsed, Tokens};
 use ruff_python_semantic::{
@@ -125,7 +125,7 @@ impl<'a> Importer<'a> {
         &self,
         import: &ImportedMembers,
         at: TextSize,
-        semantic: &SemanticModel,
+        semantic: &SemanticModel<'a>,
     ) -> Result<TypingImportEdit> {
         // Generate the modified import statement.
         let content = fix::codemods::retain_imports(
@@ -135,6 +135,39 @@ impl<'a> Importer<'a> {
             self.stylist,
         )?;
 
+        // Add the import to an existing `TYPE_CHECKING` block.
+        if let Some(block) = self.preceding_type_checking_block(at) {
+            // Add the import to the existing `TYPE_CHECKING` block.
+            let type_checking_edit =
+                if let Some(statement) = Self::type_checking_binding_statement(semantic, block) {
+                    if statement == import.statement {
+                        // Special-case: if the `TYPE_CHECKING` symbol is imported as part of the same
+                        // statement that we're modifying, avoid adding a no-op edit. For example, here,
+                        // the `TYPE_CHECKING` no-op edit would overlap with the edit to remove `Final`
+                        // from the import:
+                        // ```python
+                        // from __future__ import annotations
+                        //
+                        // from typing import Final, TYPE_CHECKING
+                        //
+                        // Const: Final[dict] = {}
+                        // ```
+                        None
+                    } else {
+                        Some(Edit::range_replacement(
+                            self.locator.slice(statement.range()).to_string(),
+                            statement.range(),
+                        ))
+                    }
+                } else {
+                    None
+                };
+            return Ok(TypingImportEdit {
+                type_checking_edit,
+                add_import_edit: self.add_to_type_checking_block(&content, block.start()),
+            });
+        }
+
         // Import the `TYPE_CHECKING` symbol from the typing module.
         let (type_checking_edit, type_checking) =
             if let Some(type_checking) = Self::find_type_checking(at, semantic)? {
@@ -179,27 +212,36 @@ impl<'a> Importer<'a> {
                 (Some(edit), name)
             };
 
-        // Add the import to a `TYPE_CHECKING` block.
-        let add_import_edit = if let Some(block) = self.preceding_type_checking_block(at) {
-            // Add the import to the `TYPE_CHECKING` block.
-            self.add_to_type_checking_block(&content, block.start())
-        } else {
-            // Add the import to a new `TYPE_CHECKING` block.
-            self.add_type_checking_block(
+        // Add the import to a new `TYPE_CHECKING` block.
+        Ok(TypingImportEdit {
+            type_checking_edit,
+            add_import_edit: self.add_type_checking_block(
                 &format!(
                     "{}if {type_checking}:{}{}",
                     self.stylist.line_ending().as_str(),
                     self.stylist.line_ending().as_str(),
                     indent(&content, self.stylist.indentation())
                 ),
                 at,
-            )?
+            )?,
+        })
+    }
+
+    fn type_checking_binding_statement(
+        semantic: &SemanticModel<'a>,
+        type_checking_block: &Stmt,
+    ) -> Option<&'a Stmt> {
+        let Stmt::If(ast::StmtIf { test, .. }) = type_checking_block else {
+            return None;
         };
 
-        Ok(TypingImportEdit {
-            type_checking_edit,
-            add_import_edit,
-        })
+        let mut source = test;
+        while let Expr::Attribute(ast::ExprAttribute { value, .. }) = source.as_ref() {
+            source = value;
+        }
+        semantic
+            .binding(semantic.resolve_name(source.as_name_expr()?)?)
+            .statement(semantic)
     }
 
     /// Find a reference to `typing.TYPE_CHECKING`.

crates/ruff_linter/src/rules/flake8_simplify/snapshots/ruff_linter__rules__flake8_simplify__tests__preview__SIM108_SIM108.py.snap

@@ -226,6 +226,33 @@ SIM108.py:167:1: SIM108 [*] Use ternary operator `z = 1 if True else other` inst
 172 169 | if False:
 173 170 |     z = 1
 
+SIM108.py:172:1: SIM108 [*] Use ternary operator `z = 1 if False else other` instead of `if`-`else`-block
+    |
+170 |       z = other
+171 |
+172 | / if False:
+173 | |     z = 1
+174 | | else:
+175 | |     z = other
+    | |_____________^ SIM108
+176 |
+177 |   if 1:
+    |
+    = help: Replace `if`-`else`-block with `z = 1 if False else other`
+
+ℹ Unsafe fix
+169 169 | else:
+170 170 |     z = other
+171 171 | 
+172     |-if False:
+173     |-    z = 1
+174     |-else:
+175     |-    z = other
+    172 |+z = 1 if False else other
+176 173 | 
+177 174 | if 1:
+178 175 |     z = True
+
 SIM108.py:177:1: SIM108 [*] Use ternary operator `z = True if 1 else other` instead of `if`-`else`-block
     |
 175 |       z = other

crates/ruff_linter/src/rules/flake8_type_checking/mod.rs

@@ -524,4 +524,41 @@ mod tests {
         );
         assert_messages!(snapshot, diagnostics);
     }
+
+    #[test_case(
+        r"
+        from __future__ import annotations
+
+        TYPE_CHECKING = False
+        if TYPE_CHECKING:
+            from types import TracebackType
+
+        def foo(tb: TracebackType): ...
+    ",
+        "github_issue_15681_regression_test"
+    )]
+    #[test_case(
+        r"
+        from __future__ import annotations
+
+        import pathlib  # TC003
+
+        TYPE_CHECKING = False
+        if TYPE_CHECKING:
+            from types import TracebackType
+
+        def foo(tb: TracebackType) -> pathlib.Path: ...
+    ",
+        "github_issue_15681_fix_test"
+    )]
+    fn contents_preview(contents: &str, snapshot: &str) {
+        let diagnostics = test_snippet(
+            contents,
+            &settings::LinterSettings {
+                preview: settings::types::PreviewMode::Enabled,
+                ..settings::LinterSettings::for_rules(Linter::Flake8TypeChecking.rules())
+            },
+        );
+        assert_messages!(snapshot, diagnostics);
+    }
 }

crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__github_issue_15681_fix_test.snap

@@ -0,0 +1,26 @@
+---
+source: crates/ruff_linter/src/rules/flake8_type_checking/mod.rs
+---
+<filename>:4:8: TC003 [*] Move standard library import `pathlib` into a type-checking block
+  |
+2 | from __future__ import annotations
+3 |
+4 | import pathlib  # TC003
+  |        ^^^^^^^ TC003
+5 |
+6 | TYPE_CHECKING = False
+  |
+  = help: Move into type-checking block
+
+ℹ Unsafe fix
+1 1 | 
+2 2 | from __future__ import annotations
+3 3 | 
+4   |-import pathlib  # TC003
+5 4 | 
+6 5 | TYPE_CHECKING = False
+7 6 | if TYPE_CHECKING:
+  7 |+    import pathlib
+8 8 |     from types import TracebackType
+9 9 | 
+10 10 | def foo(tb: TracebackType) -> pathlib.Path: ...

crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__github_issue_15681_regression_test.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff_linter/src/rules/flake8_type_checking/mod.rs
+---
+

crates/ruff_python_semantic/src/analyze/typing.rs

@@ -382,6 +382,22 @@ pub fn is_mutable_expr(expr: &Expr, semantic: &SemanticModel) -> bool {
 pub fn is_type_checking_block(stmt: &ast::StmtIf, semantic: &SemanticModel) -> bool {
     let ast::StmtIf { test, .. } = stmt;
 
+    if semantic.use_new_type_checking_block_detection_semantics() {
+        return match test.as_ref() {
+            // As long as the symbol's name is "TYPE_CHECKING" we will treat it like `typing.TYPE_CHECKING`
+            // for this specific check even if it's defined somewhere else, like the current module.
+            // Ex) `if TYPE_CHECKING:`
+            Expr::Name(ast::ExprName { id, .. }) => {
+                id == "TYPE_CHECKING"
+                // Ex) `if TC:` with `from typing import TYPE_CHECKING as TC`
+                || semantic.match_typing_expr(test, "TYPE_CHECKING")
+            }
+            // Ex) `if typing.TYPE_CHECKING:`
+            Expr::Attribute(ast::ExprAttribute { attr, .. }) => attr == "TYPE_CHECKING",
+            _ => false,
+        };
+    }
+
     // Ex) `if False:`
     if is_const_false(test) {
         return true;

crates/ruff_python_semantic/src/model.rs

@@ -2014,6 +2014,18 @@ impl<'a> SemanticModel<'a> {
             .intersects(SemanticModelFlags::DEFERRED_CLASS_BASE)
     }
 
+    /// Return `true` if we should use the new semantics to recognize
+    /// type checking blocks. Previously we only recognized type checking
+    /// blocks if `TYPE_CHECKING` was imported from a typing module.
+    ///
+    /// With this feature flag enabled we recognize any symbol named
+    /// `TYPE_CHECKING`, regardless of where it comes from to mirror
+    /// what mypy and pyright do.
+    pub const fn use_new_type_checking_block_detection_semantics(&self) -> bool {
+        self.flags
+            .intersects(SemanticModelFlags::NEW_TYPE_CHECKING_BLOCK_DETECTION)
+    }
+
     /// Return an iterator over all bindings shadowed by the given [`BindingId`], within the
     /// containing scope, and across scopes.
     pub fn shadowed_bindings(
@@ -2545,6 +2557,14 @@ bitflags! {
         /// [#13824]: https://github.com/astral-sh/ruff/issues/13824
         const NO_TYPE_CHECK = 1 << 30;
 
+        /// The model special-cases any symbol named `TYPE_CHECKING`.
+        ///
+        /// Previously we only recognized `TYPE_CHECKING` if it was part of
+        /// one of the configured `typing` modules. This flag exists to
+        /// test out the semantic change only in preview. This flag will go
+        /// away once this change has been stabilized.
+        const NEW_TYPE_CHECKING_BLOCK_DETECTION = 1 << 31;
+
         /// The context is in any type annotation.
         const ANNOTATION = Self::TYPING_ONLY_ANNOTATION.bits() | Self::RUNTIME_EVALUATED_ANNOTATION.bits() | Self::RUNTIME_REQUIRED_ANNOTATION.bits();