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,6 +47,8 @@ pub struct RuntimeImportInTypeCheckingBlock {
 }
 
 impl Violation for RuntimeImportInTypeCheckingBlock {
+    const AUTOFIX: AutofixKind = AutofixKind::Sometimes;
+
     #[derive_message_formats]
     fn message(&self) -> String {
         let RuntimeImportInTypeCheckingBlock { full_name } = self;
@@ -50,6 +56,10 @@ impl Violation for RuntimeImportInTypeCheckingBlock {
             "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,7 +1,7 @@
 ---
 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
@@ -9,5 +9,16 @@ TCH004_11.py:4:24: TCH004 Move import `typing.List` out of type-checking block.
 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,7 +1,7 @@
 ---
 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
@@ -9,5 +9,18 @@ TCH004_12.py:6:33: TCH004 Move import `collections.abc.Callable` out of type-che
 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
 
-TCH004_5.py:4:40: TCH004 Move import `typing.Set` 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 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.
   |
 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,7 +1,7 @@
 ---
 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
@@ -9,8 +9,20 @@ TCH004_9.py:4:24: TCH004 Move import `typing.Tuple` out of type-checking block.
 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
@@ -18,5 +30,17 @@ TCH004_9.py:4:31: TCH004 Move import `typing.List` out of type-checking block. I
 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 | 
 
 

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

@@ -1,15 +1,28 @@
 ---
 source: crates/ruff/src/rules/flake8_type_checking/mod.rs
 ---
-runtime_evaluated_base_classes_1.py:10:12: TCH004 Move import `datetime` out of type-checking block. Import is used for more than type hinting.
+runtime_evaluated_base_classes_1.py:10:12: TCH004 [*] Move import `datetime` out of type-checking block. Import is used for more than type hinting.
    |
 10 | if TYPE_CHECKING:
 11 |     import datetime  # TCH004
    |            ^^^^^^^^ TCH004
 12 |     from array import array  # TCH004
    |
+   = help: Move out of type-checking block
 
-runtime_evaluated_base_classes_1.py:11:23: TCH004 Move import `array.array` out of type-checking block. Import is used for more than type hinting.
+ℹ Suggested fix
+5  5  | 
+6  6  | import pydantic
+7  7  | from pydantic import BaseModel
+   8  |+import datetime
+8  9  | 
+9  10 | if TYPE_CHECKING:
+10    |-    import datetime  # TCH004
+11 11 |     from array import array  # TCH004
+12 12 | 
+13 13 |     import pandas  # TCH004
+
+runtime_evaluated_base_classes_1.py:11:23: TCH004 [*] Move import `array.array` out of type-checking block. Import is used for more than type hinting.
    |
 11 | if TYPE_CHECKING:
 12 |     import datetime  # TCH004
@@ -18,8 +31,22 @@ runtime_evaluated_base_classes_1.py:11:23: TCH004 Move import `array.array` out
 14 | 
 15 |     import pandas  # TCH004
    |
+   = help: Move out of type-checking block
 
-runtime_evaluated_base_classes_1.py:13:12: TCH004 Move import `pandas` out of type-checking block. Import is used for more than type hinting.
+ℹ Suggested fix
+5  5  | 
+6  6  | import pydantic
+7  7  | from pydantic import BaseModel
+   8  |+from array import array
+8  9  | 
+9  10 | if TYPE_CHECKING:
+10 11 |     import datetime  # TCH004
+11    |-    from array import array  # TCH004
+12 12 | 
+13 13 |     import pandas  # TCH004
+14 14 |     import pyproj
+
+runtime_evaluated_base_classes_1.py:13:12: TCH004 [*] Move import `pandas` out of type-checking block. Import is used for more than type hinting.
    |
 13 |     from array import array  # TCH004
 14 | 
@@ -27,5 +54,21 @@ runtime_evaluated_base_classes_1.py:13:12: TCH004 Move import `pandas` out of ty
    |            ^^^^^^ TCH004
 16 |     import pyproj
    |
+   = help: Move out of type-checking block
+
+ℹ Suggested fix
+5  5  | 
+6  6  | import pydantic
+7  7  | from pydantic import BaseModel
+   8  |+import pandas
+8  9  | 
+9  10 | if TYPE_CHECKING:
+10 11 |     import datetime  # TCH004
+11 12 |     from array import array  # TCH004
+12 13 | 
+13    |-    import pandas  # TCH004
+14 14 |     import pyproj
+15 15 | 
+16 16 | 
 
 

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

@@ -1,15 +1,28 @@
 ---
 source: crates/ruff/src/rules/flake8_type_checking/mod.rs
 ---
-runtime_evaluated_decorators_1.py:12:12: TCH004 Move import `datetime` out of type-checking block. Import is used for more than type hinting.
+runtime_evaluated_decorators_1.py:12:12: TCH004 [*] Move import `datetime` out of type-checking block. Import is used for more than type hinting.
    |
 12 | if TYPE_CHECKING:
 13 |     import datetime  # TCH004
    |            ^^^^^^^^ TCH004
 14 |     from array import array  # TCH004
    |
+   = help: Move out of type-checking block
 
-runtime_evaluated_decorators_1.py:13:23: TCH004 Move import `array.array` out of type-checking block. Import is used for more than type hinting.
+ℹ Suggested fix
+7  7  | from attrs import frozen
+8  8  | 
+9  9  | import numpy
+   10 |+import datetime
+10 11 | 
+11 12 | if TYPE_CHECKING:
+12    |-    import datetime  # TCH004
+13 13 |     from array import array  # TCH004
+14 14 | 
+15 15 |     import pandas  # TCH004
+
+runtime_evaluated_decorators_1.py:13:23: TCH004 [*] Move import `array.array` out of type-checking block. Import is used for more than type hinting.
    |
 13 | if TYPE_CHECKING:
 14 |     import datetime  # TCH004
@@ -18,8 +31,22 @@ runtime_evaluated_decorators_1.py:13:23: TCH004 Move import `array.array` out of
 16 | 
 17 |     import pandas  # TCH004
    |
+   = help: Move out of type-checking block
 
-runtime_evaluated_decorators_1.py:15:12: TCH004 Move import `pandas` out of type-checking block. Import is used for more than type hinting.
+ℹ Suggested fix
+7  7  | from attrs import frozen
+8  8  | 
+9  9  | import numpy
+   10 |+from array import array
+10 11 | 
+11 12 | if TYPE_CHECKING:
+12 13 |     import datetime  # TCH004
+13    |-    from array import array  # TCH004
+14 14 | 
+15 15 |     import pandas  # TCH004
+16 16 |     import pyproj
+
+runtime_evaluated_decorators_1.py:15:12: TCH004 [*] Move import `pandas` out of type-checking block. Import is used for more than type hinting.
    |
 15 |     from array import array  # TCH004
 16 | 
@@ -27,5 +54,21 @@ runtime_evaluated_decorators_1.py:15:12: TCH004 Move import `pandas` out of type
    |            ^^^^^^ TCH004
 18 |     import pyproj
    |
+   = help: Move out of type-checking block
+
+ℹ Suggested fix
+7  7  | from attrs import frozen
+8  8  | 
+9  9  | import numpy
+   10 |+import pandas
+10 11 | 
+11 12 | if TYPE_CHECKING:
+12 13 |     import datetime  # TCH004
+13 14 |     from array import array  # TCH004
+14 15 | 
+15    |-    import pandas  # TCH004
+16 16 |     import pyproj
+17 17 | 
+18 18 |