Add autofix to move runtime-imports out of type-checking blocks (#4743)

crates/ruff/src/importer/mod.rs

@@ -75,6 +75,37 @@ impl<'a> Importer<'a> {
         }
     }
 
+    /// Move an existing import to the top-level, thereby making it available at runtime.
+    ///
+    /// If there are no existing imports, the new import will be added at the top
+    /// of the file. Otherwise, it will be added after the most recent top-level
+    /// import statement.
+    pub(crate) fn runtime_import_edit(
+        &self,
+        import: &StmtImport,
+        at: TextSize,
+    ) -> Result<RuntimeImportEdit> {
+        // Generate the modified import statement.
+        let content = autofix::codemods::retain_imports(
+            &[import.full_name],
+            import.stmt,
+            self.locator,
+            self.stylist,
+        )?;
+
+        // Add the import to the top-level.
+        let insertion = if let Some(stmt) = self.preceding_import(at) {
+            // Insert after the last top-level import.
+            Insertion::end_of_statement(stmt, self.locator, self.stylist)
+        } else {
+            // Insert at the start of the file.
+            Insertion::start_of_file(self.python_ast, self.locator, self.stylist)
+        };
+        let add_import_edit = insertion.into_edit(&content);
+
+        Ok(RuntimeImportEdit { add_import_edit })
+    }
+
     /// Move an existing import into a `TYPE_CHECKING` block.
     ///
     /// If there are no existing `TYPE_CHECKING` blocks, a new one will be added at the top
@@ -344,6 +375,19 @@ impl<'a> Importer<'a> {
     }
 }
 
+/// An edit to the top-level of a module, making it available at runtime.
+#[derive(Debug)]
+pub(crate) struct RuntimeImportEdit {
+    /// The edit to add the import to the top-level of the module.
+    add_import_edit: Edit,
+}
+
+impl RuntimeImportEdit {
+    pub(crate) fn into_edits(self) -> Vec<Edit> {
+        vec![self.add_import_edit]
+    }
+}
+
 /// An edit to an import to a typing-only context.
 #[derive(Debug)]
 pub(crate) struct TypingImportEdit {

crates/ruff/src/rules/flake8_type_checking/rules/runtime_import_in_type_checking_block.rs

@@ -1,10 +1,14 @@
-use ruff_diagnostics::{Diagnostic, Violation};
+use rustpython_parser::ast::Stmt;
+
+use ruff_diagnostics::{AutofixKind, Diagnostic, Fix, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_semantic::binding::{
     Binding, BindingKind, FromImportation, Importation, SubmoduleImportation,
 };
 
+use crate::autofix;
 use crate::checkers::ast::Checker;
+use crate::importer::StmtImport;
 use crate::registry::AsRule;
 
 /// ## What it does
@@ -43,13 +47,19 @@ pub struct RuntimeImportInTypeCheckingBlock {
 }
 
 impl Violation for RuntimeImportInTypeCheckingBlock {
+    const AUTOFIX: AutofixKind = AutofixKind::Sometimes;
+
     #[derive_message_formats]
     fn message(&self) -> String {
         let RuntimeImportInTypeCheckingBlock { full_name } = self;
         format!(
             "Move import `{full_name}` out of type-checking block. Import is used for more than type hinting."
         )
     }
+
+    fn autofix_title(&self) -> Option<String> {
+        Some("Move out of type-checking block".to_string())
+    }
 }
 
 /// TCH004
@@ -65,6 +75,10 @@ pub(crate) fn runtime_import_in_type_checking_block(
         _ => return,
     };
 
+    let Some(reference_id) = binding.references.first() else {
+        return;
+    };
+
     if binding.context.is_typing()
         && binding.references().any(|reference_id| {
             checker
@@ -75,12 +89,49 @@ pub(crate) fn runtime_import_in_type_checking_block(
                 .is_runtime()
         })
     {
-        let diagnostic = Diagnostic::new(
+        let mut diagnostic = Diagnostic::new(
             RuntimeImportInTypeCheckingBlock {
                 full_name: full_name.to_string(),
             },
             binding.range,
         );
+
+        if checker.patch(diagnostic.kind.rule()) {
+            diagnostic.try_set_fix(|| {
+                // Step 1) Remove the import from the type-checking block.
+                // SAFETY: All non-builtin bindings have a source.
+                let source = binding.source.unwrap();
+                let deleted: Vec<&Stmt> = checker.deletions.iter().map(Into::into).collect();
+                let stmt = checker.semantic_model().stmts[source];
+                let parent = checker
+                    .semantic_model()
+                    .stmts
+                    .parent_id(source)
+                    .map(|id| checker.semantic_model().stmts[id]);
+                let remove_import_edit = autofix::edits::remove_unused_imports(
+                    std::iter::once(full_name),
+                    stmt,
+                    parent,
+                    &deleted,
+                    checker.locator,
+                    checker.indexer,
+                    checker.stylist,
+                )?;
+
+                // Step 2) Add the import to the top-level.
+                let reference = checker.semantic_model().references.resolve(*reference_id);
+                let add_import_edit = checker.importer.runtime_import_edit(
+                    &StmtImport { stmt, full_name },
+                    reference.range().start(),
+                )?;
+
+                Ok(Fix::suggested_edits(
+                    remove_import_edit,
+                    add_import_edit.into_edits(),
+                ))
+            });
+        }
+
         if checker.enabled(diagnostic.kind.rule()) {
             diagnostics.push(diagnostic);
         }

crates/ruff/src/rules/flake8_type_checking/snapshots/ruff__rules__flake8_type_checking__tests__runtime-import-in-type-checking-block_TCH004_1.py.snap

@@ -1,12 +1,22 @@
 ---
 source: crates/ruff/src/rules/flake8_type_checking/mod.rs
 ---
-TCH004_1.py:4:26: TCH004 Move import `datetime.datetime` out of type-checking block. Import is used for more than type hinting.
+TCH004_1.py:4:26: TCH004 [*] Move import `datetime.datetime` out of type-checking block. Import is used for more than type hinting.
   |
 4 | if TYPE_CHECKING:
 5 |     from datetime import datetime
   |                          ^^^^^^^^ TCH004
 6 | x = datetime
   |
+  = help: Move out of type-checking block
+
+ℹ Suggested fix
+1 1 | from typing import TYPE_CHECKING
+  2 |+from datetime import datetime
+2 3 | 
+3 4 | if TYPE_CHECKING:
+4   |-    from datetime import datetime
+  5 |+    pass
+5 6 | x = datetime
 
 

crates/ruff/src/rules/flake8_type_checking/snapshots/ruff__rules__flake8_type_checking__tests__runtime-import-in-type-checking-block_TCH004_11.py.snap

@@ -1,13 +1,24 @@
 ---
 source: crates/ruff/src/rules/flake8_type_checking/mod.rs
 ---
-TCH004_11.py:4:24: TCH004 Move import `typing.List` out of type-checking block. Import is used for more than type hinting.
+TCH004_11.py:4:24: TCH004 [*] Move import `typing.List` out of type-checking block. Import is used for more than type hinting.
   |
 4 | if TYPE_CHECKING:
 5 |     from typing import List
   |                        ^^^^ TCH004
 6 | 
 7 | __all__ = ("List",)
   |
+  = help: Move out of type-checking block
+
+ℹ Suggested fix
+1 1 | from typing import TYPE_CHECKING
+  2 |+from typing import List
+2 3 | 
+3 4 | if TYPE_CHECKING:
+4   |-    from typing import List
+  5 |+    pass
+5 6 | 
+6 7 | __all__ = ("List",)
 
 

crates/ruff/src/rules/flake8_type_checking/snapshots/ruff__rules__flake8_type_checking__tests__runtime-import-in-type-checking-block_TCH004_12.py.snap

@@ -1,13 +1,26 @@
 ---
 source: crates/ruff/src/rules/flake8_type_checking/mod.rs
 ---
-TCH004_12.py:6:33: TCH004 Move import `collections.abc.Callable` out of type-checking block. Import is used for more than type hinting.
+TCH004_12.py:6:33: TCH004 [*] Move import `collections.abc.Callable` out of type-checking block. Import is used for more than type hinting.
   |
 6 | if TYPE_CHECKING:
 7 |     from collections.abc import Callable
   |                                 ^^^^^^^^ TCH004
 8 | 
 9 | AnyCallable: TypeAlias = Callable[..., Any]
   |
+  = help: Move out of type-checking block
+
+ℹ Suggested fix
+1 1 | from __future__ import annotations
+2 2 | 
+3 3 | from typing import Any, TYPE_CHECKING, TypeAlias
+  4 |+from collections.abc import Callable
+4 5 | 
+5 6 | if TYPE_CHECKING:
+6   |-    from collections.abc import Callable
+  7 |+    pass
+7 8 | 
+8 9 | AnyCallable: TypeAlias = Callable[..., Any]
 
 

crates/ruff/src/rules/flake8_type_checking/snapshots/ruff__rules__flake8_type_checking__tests__runtime-import-in-type-checking-block_TCH004_2.py.snap

@@ -1,11 +1,23 @@
 ---
 source: crates/ruff/src/rules/flake8_type_checking/mod.rs
 ---
-TCH004_2.py:4:26: TCH004 Move import `datetime.date` out of type-checking block. Import is used for more than type hinting.
+TCH004_2.py:4:26: TCH004 [*] Move import `datetime.date` out of type-checking block. Import is used for more than type hinting.
   |
 4 | if TYPE_CHECKING:
 5 |     from datetime import date
   |                          ^^^^ TCH004
   |
+  = help: Move out of type-checking block
+
+ℹ Suggested fix
+1 1 | from typing import TYPE_CHECKING
+  2 |+from datetime import date
+2 3 | 
+3 4 | if TYPE_CHECKING:
+4   |-    from datetime import date
+  5 |+    pass
+5 6 | 
+6 7 | 
+7 8 | def example():
 
 

crates/ruff/src/rules/flake8_type_checking/snapshots/ruff__rules__flake8_type_checking__tests__runtime-import-in-type-checking-block_TCH004_4.py.snap

@@ -1,11 +1,23 @@
 ---
 source: crates/ruff/src/rules/flake8_type_checking/mod.rs
 ---
-TCH004_4.py:4:24: TCH004 Move import `typing.Any` out of type-checking block. Import is used for more than type hinting.
+TCH004_4.py:4:24: TCH004 [*] Move import `typing.Any` out of type-checking block. Import is used for more than type hinting.
   |
 4 | if TYPE_CHECKING:
 5 |     from typing import Any
   |                        ^^^ TCH004
   |
+  = help: Move out of type-checking block
+
+ℹ Suggested fix
+1 1 | from typing import TYPE_CHECKING, Type
+  2 |+from typing import Any
+2 3 | 
+3 4 | if TYPE_CHECKING:
+4   |-    from typing import Any
+  5 |+    pass
+5 6 | 
+6 7 | 
+7 8 | def example(*args: Any, **kwargs: Any):
 
 

crates/ruff/src/rules/flake8_type_checking/snapshots/ruff__rules__flake8_type_checking__tests__runtime-import-in-type-checking-block_TCH004_5.py.snap

@@ -1,25 +1,61 @@
 ---
 source: crates/ruff/src/rules/flake8_type_checking/mod.rs
 ---
-TCH004_5.py:4:24: TCH004 Move import `typing.List` out of type-checking block. Import is used for more than type hinting.
+TCH004_5.py:4:24: TCH004 [*] Move import `typing.List` out of type-checking block. Import is used for more than type hinting.
   |
 4 | if TYPE_CHECKING:
 5 |     from typing import List, Sequence, Set
   |                        ^^^^ TCH004
   |
+  = help: Move out of type-checking block
 
-TCH004_5.py:4:30: TCH004 Move import `typing.Sequence` out of type-checking block. Import is used for more than type hinting.
+ℹ Suggested fix
+1 1 | from typing import TYPE_CHECKING
+  2 |+from typing import List
+2 3 | 
+3 4 | if TYPE_CHECKING:
+4   |-    from typing import List, Sequence, Set
+  5 |+    from typing import Sequence, Set
+5 6 | 
+6 7 | 
+7 8 | def example(a: List[int], /, b: Sequence[int], *, c: Set[int]):
+
+TCH004_5.py:4:30: TCH004 [*] Move import `typing.Sequence` out of type-checking block. Import is used for more than type hinting.
   |
 4 | if TYPE_CHECKING:
 5 |     from typing import List, Sequence, Set
   |                              ^^^^^^^^ TCH004
   |
+  = help: Move out of type-checking block
+
+ℹ Suggested fix
+1 1 | from typing import TYPE_CHECKING
+  2 |+from typing import Sequence
+2 3 | 
+3 4 | if TYPE_CHECKING:
+4   |-    from typing import List, Sequence, Set
+  5 |+    from typing import List, Set
+5 6 | 
+6 7 | 
+7 8 | def example(a: List[int], /, b: Sequence[int], *, c: Set[int]):
 
-TCH004_5.py:4:40: TCH004 Move import `typing.Set` out of type-checking block. Import is used for more than type hinting.
+TCH004_5.py:4:40: TCH004 [*] Move import `typing.Set` out of type-checking block. Import is used for more than type hinting.
   |
 4 | if TYPE_CHECKING:
 5 |     from typing import List, Sequence, Set
   |                                        ^^^ TCH004
   |
+  = help: Move out of type-checking block
+
+ℹ Suggested fix
+1 1 | from typing import TYPE_CHECKING
+  2 |+from typing import Set
+2 3 | 
+3 4 | if TYPE_CHECKING:
+4   |-    from typing import List, Sequence, Set
+  5 |+    from typing import List, Sequence
+5 6 | 
+6 7 | 
+7 8 | def example(a: List[int], /, b: Sequence[int], *, c: Set[int]):
 
 

crates/ruff/src/rules/flake8_type_checking/snapshots/ruff__rules__flake8_type_checking__tests__runtime-import-in-type-checking-block_TCH004_9.py.snap

@@ -1,22 +1,46 @@
 ---
 source: crates/ruff/src/rules/flake8_type_checking/mod.rs
 ---
-TCH004_9.py:4:24: TCH004 Move import `typing.Tuple` out of type-checking block. Import is used for more than type hinting.
+TCH004_9.py:4:24: TCH004 [*] Move import `typing.Tuple` out of type-checking block. Import is used for more than type hinting.
   |
 4 | if TYPE_CHECKING:
 5 |     from typing import Tuple, List, Dict
   |                        ^^^^^ TCH004
 6 | 
 7 | x: Tuple
   |
+  = help: Move out of type-checking block
 
-TCH004_9.py:4:31: TCH004 Move import `typing.List` out of type-checking block. Import is used for more than type hinting.
+ℹ Suggested fix
+1 1 | from typing import TYPE_CHECKING
+  2 |+from typing import Tuple
+2 3 | 
+3 4 | if TYPE_CHECKING:
+4   |-    from typing import Tuple, List, Dict
+  5 |+    from typing import List, Dict
+5 6 | 
+6 7 | x: Tuple
+7 8 | 
+
+TCH004_9.py:4:31: TCH004 [*] Move import `typing.List` out of type-checking block. Import is used for more than type hinting.
   |
 4 | if TYPE_CHECKING:
 5 |     from typing import Tuple, List, Dict
   |                               ^^^^ TCH004
 6 | 
 7 | x: Tuple
   |
+  = help: Move out of type-checking block
+
+ℹ Suggested fix
+1 1 | from typing import TYPE_CHECKING
+  2 |+from typing import List
+2 3 | 
+3 4 | if TYPE_CHECKING:
+4   |-    from typing import Tuple, List, Dict
+  5 |+    from typing import Tuple, Dict
+5 6 | 
+6 7 | x: Tuple
+7 8 |