Improve detection of TYPE_CHECKING blocks imported from typing_extensions or _typeshed

crates/ruff_linter/resources/test/fixtures/flake8_type_checking/TCH002.py

@@ -172,3 +172,14 @@ def f():
     from module import Member
 
     x: Member = 1
+
+
+def f():
+    from typing_extensions import TYPE_CHECKING
+
+    from pandas import y
+
+    if TYPE_CHECKING:
+        _type = x
+    elif True:
+        _type = y

crates/ruff_linter/resources/test/fixtures/flake8_type_checking/TCH004_16.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from typing_extensions import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pandas import DataFrame
+
+
+def example() -> DataFrame:
+    pass

crates/ruff_linter/resources/test/fixtures/flake8_type_checking/TCH004_17.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from typing_extensions import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pandas import DataFrame
+
+
+def example() -> DataFrame:
+    x = DataFrame()

crates/ruff_linter/resources/test/fixtures/flake8_type_checking/TCH005.py

@@ -37,3 +37,9 @@ class Test:
 
 if 0:
     x: List
+
+
+from typing_extensions import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    pass  # TCH005

crates/ruff_linter/resources/test/fixtures/flake8_type_checking/typing_modules_1.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from typing_extensions import Self
+
+
+def func():
+    from pandas import DataFrame
+
+    df: DataFrame

crates/ruff_linter/resources/test/fixtures/flake8_type_checking/typing_modules_2.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+import typing_extensions
+
+
+def func():
+    from pandas import DataFrame
+
+    df: DataFrame

crates/ruff_linter/src/importer/mod.rs

@@ -132,11 +132,7 @@ impl<'a> Importer<'a> {
         )?;
 
         // Import the `TYPE_CHECKING` symbol from the typing module.
-        let (type_checking_edit, type_checking) = self.get_or_import_symbol(
-            &ImportRequest::import_from("typing", "TYPE_CHECKING"),
-            at,
-            semantic,
-        )?;
+        let (type_checking_edit, type_checking) = self.get_or_import_type_checking(at, semantic)?;
 
         // Add the import to a `TYPE_CHECKING` block.
         let add_import_edit = if let Some(block) = self.preceding_type_checking_block(at) {
@@ -161,6 +157,30 @@ impl<'a> Importer<'a> {
         })
     }
 
+    /// Generate an [`Edit`] to reference `typing.TYPE_CHECKING`. Returns the [`Edit`] necessary to
+    /// make the symbol available in the current scope along with the bound name of the symbol.
+    fn get_or_import_type_checking(
+        &self,
+        at: TextSize,
+        semantic: &SemanticModel,
+    ) -> Result<(Edit, String), ResolutionError> {
+        for module in semantic.typing_modules() {
+            if let Some((edit, name)) = self.get_symbol(
+                &ImportRequest::import_from(module, "TYPE_CHECKING"),
+                at,
+                semantic,
+            )? {
+                return Ok((edit, name));
+            }
+        }
+
+        self.import_symbol(
+            &ImportRequest::import_from("typing", "TYPE_CHECKING"),
+            at,
+            semantic,
+        )
+    }
+
     /// Generate an [`Edit`] to reference the given symbol. Returns the [`Edit`] necessary to make
     /// the symbol available in the current scope along with the bound name of the symbol.
     ///

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

@@ -23,6 +23,8 @@ mod tests {
     #[test_case(Rule::RuntimeImportInTypeCheckingBlock, Path::new("TCH004_13.py"))]
     #[test_case(Rule::RuntimeImportInTypeCheckingBlock, Path::new("TCH004_14.pyi"))]
     #[test_case(Rule::RuntimeImportInTypeCheckingBlock, Path::new("TCH004_15.py"))]
+    #[test_case(Rule::RuntimeImportInTypeCheckingBlock, Path::new("TCH004_16.py"))]
+    #[test_case(Rule::RuntimeImportInTypeCheckingBlock, Path::new("TCH004_17.py"))]
     #[test_case(Rule::RuntimeImportInTypeCheckingBlock, Path::new("TCH004_2.py"))]
     #[test_case(Rule::RuntimeImportInTypeCheckingBlock, Path::new("TCH004_3.py"))]
     #[test_case(Rule::RuntimeImportInTypeCheckingBlock, Path::new("TCH004_4.py"))]
@@ -36,6 +38,8 @@ mod tests {
     #[test_case(Rule::TypingOnlyStandardLibraryImport, Path::new("snapshot.py"))]
     #[test_case(Rule::TypingOnlyThirdPartyImport, Path::new("TCH002.py"))]
     #[test_case(Rule::TypingOnlyThirdPartyImport, Path::new("strict.py"))]
+    #[test_case(Rule::TypingOnlyThirdPartyImport, Path::new("typing_modules_1.py"))]
+    #[test_case(Rule::TypingOnlyThirdPartyImport, Path::new("typing_modules_2.py"))]
     fn rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.as_ref(), path.to_string_lossy());
         let diagnostics = test_path(

crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__empty-type-checking-block_TCH005.py.snap

@@ -96,4 +96,19 @@ TCH005.py:22:9: TCH005 [*] Found empty type-checking block
 24 22 | 
 25 23 | 
 
+TCH005.py:45:5: TCH005 [*] Found empty type-checking block
+   |
+44 | if TYPE_CHECKING:
+45 |     pass  # TCH005
+   |     ^^^^ TCH005
+   |
+   = help: Delete empty type-checking block
+
+ℹ Safe fix
+41 41 | 
+42 42 | from typing_extensions import TYPE_CHECKING
+43 43 | 
+44    |-if TYPE_CHECKING:
+45    |-    pass  # TCH005
+
 

crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__runtime-import-in-type-checking-block_TCH004_16.py.snap

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

crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__runtime-import-in-type-checking-block_TCH004_17.py.snap

@@ -0,0 +1,25 @@
+---
+source: crates/ruff_linter/src/rules/flake8_type_checking/mod.rs
+---
+TCH004_17.py:6:24: TCH004 [*] Move import `pandas.DataFrame` out of type-checking block. Import is used for more than type hinting.
+  |
+5 | if TYPE_CHECKING:
+6 |     from pandas import DataFrame
+  |                        ^^^^^^^^^ TCH004
+  |
+  = help: Move out of type-checking block
+
+ℹ Unsafe fix
+1 1 | from __future__ import annotations
+2 2 | 
+3 3 | from typing_extensions import TYPE_CHECKING
+  4 |+from pandas import DataFrame
+4 5 | 
+5 6 | if TYPE_CHECKING:
+6   |-    from pandas import DataFrame
+  7 |+    pass
+7 8 | 
+8 9 | 
+9 10 | def example() -> DataFrame:
+
+

crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__typing-only-third-party-import_TCH002.py.snap

@@ -248,5 +248,6 @@ TCH002.py:172:24: TCH002 [*] Move third-party import `module.Member` into a type
 172     |-    from module import Member
 173 176 | 
 174 177 |     x: Member = 1
+175 178 | 
 
 

crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__typing-only-third-party-import_typing_modules_1.py.snap

@@ -0,0 +1,29 @@
+---
+source: crates/ruff_linter/src/rules/flake8_type_checking/mod.rs
+---
+typing_modules_1.py:7:24: TCH002 [*] Move third-party import `pandas.DataFrame` into a type-checking block
+  |
+6 | def func():
+7 |     from pandas import DataFrame
+  |                        ^^^^^^^^^ TCH002
+8 | 
+9 |     df: DataFrame
+  |
+  = help: Move into type-checking block
+
+ℹ Unsafe fix
+1  1  | from __future__ import annotations
+2  2  | 
+3  3  | from typing_extensions import Self
+   4  |+from typing import TYPE_CHECKING
+   5  |+
+   6  |+if TYPE_CHECKING:
+   7  |+    from pandas import DataFrame
+4  8  | 
+5  9  | 
+6  10 | def func():
+7     |-    from pandas import DataFrame
+8  11 | 
+9  12 |     df: DataFrame
+
+

crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__typing-only-third-party-import_typing_modules_2.py.snap

@@ -0,0 +1,27 @@
+---
+source: crates/ruff_linter/src/rules/flake8_type_checking/mod.rs
+---
+typing_modules_2.py:7:24: TCH002 [*] Move third-party import `pandas.DataFrame` into a type-checking block
+  |
+6 | def func():
+7 |     from pandas import DataFrame
+  |                        ^^^^^^^^^ TCH002
+8 | 
+9 |     df: DataFrame
+  |
+  = help: Move into type-checking block
+
+ℹ Unsafe fix
+2  2  | 
+3  3  | import typing_extensions
+4  4  | 
+   5  |+if typing_extensions.TYPE_CHECKING:
+   6  |+    from pandas import DataFrame
+   7  |+
+5  8  | 
+6  9  | def func():
+7     |-    from pandas import DataFrame
+8  10 | 
+9  11 |     df: DataFrame
+
+

crates/ruff_python_semantic/src/analyze/typing.rs

@@ -315,10 +315,7 @@ pub fn is_type_checking_block(stmt: &ast::StmtIf, semantic: &SemanticModel) -> b
     }
 
     // Ex) `if typing.TYPE_CHECKING:`
-    if semantic
-        .resolve_call_path(test)
-        .is_some_and(|call_path| matches!(call_path.as_slice(), ["typing", "TYPE_CHECKING"]))
-    {
+    if semantic.match_typing_expr(test, "TYPE_CHECKING") {
         return true;
     }
 

crates/ruff_python_semantic/src/model.rs

@@ -200,6 +200,14 @@ impl<'a> SemanticModel<'a> {
         false
     }
 
+    /// Return an iterator over the set of `typing` modules allowed in the semantic model.
+    pub fn typing_modules(&self) -> impl Iterator<Item = &str> {
+        ["typing", "_typeshed", "typing_extensions"]
+            .iter()
+            .copied()
+            .chain(self.typing_modules.iter().map(String::as_str))
+    }
+
     /// Create a new [`Binding`] for a builtin.
     pub fn push_builtin(&mut self) -> BindingId {
         self.bindings.push(Binding {