Recognize all symbols named TYPE_CHECKING for in_type_checking_block

crates/ruff_linter/src/importer/mod.rs

@@ -135,7 +135,20 @@ impl<'a> Importer<'a> {
             self.stylist,
         )?;
 
+        // Add the import to a `TYPE_CHECKING` block.
+        if let Some(block) = self.preceding_type_checking_block(at) {
+            // Add the import to the existing `TYPE_CHECKING` block.
+            return Ok(TypingImportEdit {
+                type_checking_edit: None,
+                add_import_edit: self.add_to_type_checking_block(&content, block.start()),
+            });
+        }
+
         // Import the `TYPE_CHECKING` symbol from the typing module.
+        // TODO: Should we provide an option to avoid this import?
+        //       E.g. either through an explicit setting, or implicitly
+        //       when `typing` isn't part of the exempt modules and there
+        //       are no other existing runtime imports of `typing`.
         let (type_checking_edit, type_checking) =
             if let Some(type_checking) = Self::find_type_checking(at, semantic)? {
                 // Special-case: if the `TYPE_CHECKING` symbol is imported as part of the same
@@ -179,26 +192,18 @@ 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,
-            )?
-        };
-
-        Ok(TypingImportEdit {
-            type_checking_edit,
-            add_import_edit,
+            )?,
         })
     }
 

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

@@ -517,6 +517,32 @@ mod tests {
     ",
         "tc010_precedence_over_tc008"
     )]
+    #[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(contents: &str, snapshot: &str) {
         let diagnostics = test_snippet(
             contents,

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

@@ -328,28 +328,31 @@ 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;
 
+    // TODO: This form is no longer supported by mypy/pyright, do we still support it?
     // Ex) `if False:`
     if is_const_false(test) {
         return true;
     }
 
-    // Ex) `if 0:`
-    if matches!(
-        test.as_ref(),
+    match test.as_ref() {
+        // TODO: This is also no longer supported
+        // Ex) `if 0:`
         Expr::NumberLiteral(ast::ExprNumberLiteral {
             value: ast::Number::Int(Int::ZERO),
             ..
-        })
-    ) {
-        return true;
-    }
-
-    // Ex) `if typing.TYPE_CHECKING:`
-    if semantic.match_typing_expr(test, "TYPE_CHECKING") {
-        return true;
+        }) => true,
+        // 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.as_str() == "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.as_str() == "TYPE_CHECKING",
+        _ => false,
     }
-
-    false
 }
 
 /// Returns `true` if the [`ast::StmtIf`] is a version-checking block (e.g., `if sys.version_info >= ...:`).