diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/TCH002.py b/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/TCH002.py index 9248c10775302..d3c36312eea7a 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/TCH002.py +++ b/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 diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/TCH004_16.py b/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/TCH004_16.py new file mode 100644 index 0000000000000..5246360f3fc65 --- /dev/null +++ b/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 diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/TCH004_17.py b/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/TCH004_17.py new file mode 100644 index 0000000000000..edcd878addf22 --- /dev/null +++ b/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() diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/TCH005.py b/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/TCH005.py index 2036166c7e231..ee1c43e0ae28f 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/TCH005.py +++ b/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 diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/typing_modules_1.py b/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/typing_modules_1.py new file mode 100644 index 0000000000000..f73487e3eb9ec --- /dev/null +++ b/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 diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/typing_modules_2.py b/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/typing_modules_2.py new file mode 100644 index 0000000000000..56db40970f087 --- /dev/null +++ b/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 diff --git a/crates/ruff_linter/src/importer/mod.rs b/crates/ruff_linter/src/importer/mod.rs index 4624f12f688e0..ecf3f38399f59 100644 --- a/crates/ruff_linter/src/importer/mod.rs +++ b/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. /// diff --git a/crates/ruff_linter/src/rules/flake8_type_checking/mod.rs b/crates/ruff_linter/src/rules/flake8_type_checking/mod.rs index cbea3cfacf60f..14bb56bdba7dd 100644 --- a/crates/ruff_linter/src/rules/flake8_type_checking/mod.rs +++ b/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( diff --git a/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__empty-type-checking-block_TCH005.py.snap b/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__empty-type-checking-block_TCH005.py.snap index e10dc87663bde..3a99c690d360b 100644 --- a/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__empty-type-checking-block_TCH005.py.snap +++ b/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 + diff --git a/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 b/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 new file mode 100644 index 0000000000000..6c5ead27428ce --- /dev/null +++ b/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 +--- + diff --git a/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 b/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 new file mode 100644 index 0000000000000..785c4c1d2e6c4 --- /dev/null +++ b/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: + + diff --git a/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__typing-only-third-party-import_TCH002.py.snap b/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__typing-only-third-party-import_TCH002.py.snap index 4221412f21d1a..7d8365daa6ef9 100644 --- a/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__typing-only-third-party-import_TCH002.py.snap +++ b/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 | diff --git a/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 b/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 new file mode 100644 index 0000000000000..3e758fc6c110d --- /dev/null +++ b/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 + + diff --git a/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 b/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 new file mode 100644 index 0000000000000..820be3f860f27 --- /dev/null +++ b/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 + + diff --git a/crates/ruff_python_semantic/src/analyze/typing.rs b/crates/ruff_python_semantic/src/analyze/typing.rs index d6032573dc433..bd18f70ba8aae 100644 --- a/crates/ruff_python_semantic/src/analyze/typing.rs +++ b/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; } diff --git a/crates/ruff_python_semantic/src/model.rs b/crates/ruff_python_semantic/src/model.rs index 28b09e653fd0c..42617ddaf7396 100644 --- a/crates/ruff_python_semantic/src/model.rs +++ b/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 { + ["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 {