Avoid infinite loop for required imports with isort: off

crates/ruff/resources/test/fixtures/isort/required_imports/off.py

@@ -0,0 +1,4 @@
+# isort: off
+
+x = 1
+# isort: on

crates/ruff/src/checkers/imports.rs

@@ -14,7 +14,7 @@ use ruff_python_stdlib::path::is_python_stub_file;
 use crate::directives::IsortDirectives;
 use crate::registry::Rule;
 use crate::rules::isort;
-use crate::rules::isort::track::{Block, ImportTracker};
+use crate::rules::isort::block::{Block, BlockBuilder};
 use crate::settings::Settings;
 
 fn extract_import_map(path: &Path, package: Option<&Path>, blocks: &[&Block]) -> Option<ImportMap> {
@@ -86,9 +86,9 @@ pub(crate) fn check_imports(
 ) -> (Vec<Diagnostic>, Option<ImportMap>) {
     let is_stub = is_python_stub_file(path);
 
-    // Extract all imports from the AST.
+    // Extract all import blocks from the AST.
     let tracker = {
-        let mut tracker = ImportTracker::new(locator, directives, is_stub);
+        let mut tracker = BlockBuilder::new(locator, directives, is_stub);
         tracker.visit_body(python_ast);
         tracker
     };
@@ -109,7 +109,7 @@ pub(crate) fn check_imports(
     }
     if settings.rules.enabled(Rule::MissingRequiredImport) {
         diagnostics.extend(isort::rules::add_required_imports(
-            &blocks, python_ast, locator, stylist, settings, is_stub,
+            python_ast, locator, stylist, settings, is_stub,
         ));
     }
 

crates/ruff/src/rules/isort/track.rs → crates/ruff/src/rules/isort/block.rs

@@ -7,21 +7,24 @@ use ruff_python_ast::statement_visitor::StatementVisitor;
 use crate::directives::IsortDirectives;
 use crate::rules::isort::helpers;
 
-#[derive(Debug, Copy, Clone)]
-pub(crate) enum Trailer {
-    Sibling,
-    ClassDef,
-    FunctionDef,
-}
-
+/// A block of imports within a Python module.
 #[derive(Debug, Default)]
 pub(crate) struct Block<'a> {
     pub(crate) nested: bool,
     pub(crate) imports: Vec<&'a Stmt>,
     pub(crate) trailer: Option<Trailer>,
 }
 
-pub(crate) struct ImportTracker<'a> {
+/// The type of trailer that should follow an import block.
+#[derive(Debug, Copy, Clone)]
+pub(crate) enum Trailer {
+    Sibling,
+    ClassDef,
+    FunctionDef,
+}
+
+/// A builder for identifying and constructing import blocks within a Python module.
+pub(crate) struct BlockBuilder<'a> {
     locator: &'a Locator<'a>,
     is_stub: bool,
     blocks: Vec<Block<'a>>,
@@ -30,7 +33,7 @@ pub(crate) struct ImportTracker<'a> {
     nested: bool,
 }
 
-impl<'a> ImportTracker<'a> {
+impl<'a> BlockBuilder<'a> {
     pub(crate) fn new(
         locator: &'a Locator<'a>,
         directives: &'a IsortDirectives,
@@ -111,7 +114,7 @@ impl<'a> ImportTracker<'a> {
     }
 }
 
-impl<'a, 'b> StatementVisitor<'b> for ImportTracker<'a>
+impl<'a, 'b> StatementVisitor<'b> for BlockBuilder<'a>
 where
     'b: 'a,
 {

crates/ruff/src/rules/isort/mod.rs

@@ -4,6 +4,7 @@ use std::collections::BTreeSet;
 use std::path::{Path, PathBuf};
 
 use annotate::annotate_imports;
+use block::{Block, Trailer};
 pub(crate) use categorize::categorize;
 use categorize::categorize_imports;
 pub use categorize::{ImportSection, ImportType};
@@ -13,7 +14,6 @@ use order::order_imports;
 use ruff_python_ast::source_code::{Locator, Stylist};
 use settings::RelativeImportsOrder;
 use sorting::cmp_either_import;
-use track::{Block, Trailer};
 use types::EitherImport::{Import, ImportFrom};
 use types::{AliasData, EitherImport, TrailingComma};
 
@@ -22,6 +22,7 @@ use crate::rules::isort::types::ImportBlock;
 use crate::settings::types::PythonVersion;
 
 mod annotate;
+pub(crate) mod block;
 mod categorize;
 mod comments;
 mod format;
@@ -32,7 +33,6 @@ pub(crate) mod rules;
 pub mod settings;
 mod sorting;
 mod split;
-pub(crate) mod track;
 mod types;
 
 #[derive(Debug)]
@@ -687,11 +687,16 @@ mod tests {
         Ok(())
     }
 
+    #[test_case(Path::new("comment.py"))]
     #[test_case(Path::new("docstring.py"))]
     #[test_case(Path::new("docstring.pyi"))]
     #[test_case(Path::new("docstring_only.py"))]
-    #[test_case(Path::new("multiline_docstring.py"))]
+    #[test_case(Path::new("docstring_with_continuation.py"))]
+    #[test_case(Path::new("docstring_with_semicolon.py"))]
     #[test_case(Path::new("empty.py"))]
+    #[test_case(Path::new("existing_import.py"))]
+    #[test_case(Path::new("multiline_docstring.py"))]
+    #[test_case(Path::new("off.py"))]
     fn required_import(path: &Path) -> Result<()> {
         let snapshot = format!("required_import_{}", path.to_string_lossy());
         let diagnostics = test_path(

crates/ruff/src/rules/isort/rules/add_required_imports.rs

@@ -11,7 +11,6 @@ use ruff_python_ast::source_code::{Locator, Stylist};
 
 use crate::importer::Importer;
 use crate::registry::Rule;
-use crate::rules::isort::track::Block;
 use crate::settings::Settings;
 
 /// ## What it does
@@ -53,58 +52,48 @@ impl AlwaysAutofixableViolation for MissingRequiredImport {
     }
 }
 
-fn contains(block: &Block, required_import: &AnyImport) -> bool {
-    block.imports.iter().any(|import| match required_import {
-        AnyImport::Import(required_import) => {
+/// Return `true` if the [`Stmt`] includes the given [`AnyImport`].
+fn includes_import(stmt: &Stmt, target: &AnyImport) -> bool {
+    match target {
+        AnyImport::Import(target) => {
             let Stmt::Import(ast::StmtImport {
                 names,
                  range: _,
-            }) = &import else {
+            }) = &stmt else {
                 return false;
             };
             names.iter().any(|alias| {
-                &alias.name == required_import.name.name
-                    && alias.asname.as_deref() == required_import.name.as_name
+                &alias.name == target.name.name && alias.asname.as_deref() == target.name.as_name
             })
         }
-        AnyImport::ImportFrom(required_import) => {
+        AnyImport::ImportFrom(target) => {
             let Stmt::ImportFrom(ast::StmtImportFrom {
                 module,
                 names,
                 level,
                  range: _,
-            }) = &import else {
+            }) = &stmt else {
                 return false;
             };
-            module.as_deref() == required_import.module
-                && level.map(|level| level.to_u32()) == required_import.level
+            module.as_deref() == target.module
+                && level.map(|level| level.to_u32()) == target.level
                 && names.iter().any(|alias| {
-                    &alias.name == required_import.name.name
-                        && alias.asname.as_deref() == required_import.name.as_name
+                    &alias.name == target.name.name
+                        && alias.asname.as_deref() == target.name.as_name
                 })
         }
-    })
+    }
 }
 
 #[allow(clippy::too_many_arguments)]
 fn add_required_import(
     required_import: &AnyImport,
-    blocks: &[&Block],
     python_ast: &Suite,
     locator: &Locator,
     stylist: &Stylist,
     settings: &Settings,
     is_stub: bool,
 ) -> Option<Diagnostic> {
-    // If the import is already present in a top-level block, don't add it.
-    if blocks
-        .iter()
-        .filter(|block| !block.nested)
-        .any(|block| contains(block, required_import))
-    {
-        return None;
-    }
-
     // Don't add imports to semantically-empty files.
     if python_ast.iter().all(is_docstring_stmt) {
         return None;
@@ -115,6 +104,14 @@ fn add_required_import(
         return None;
     }
 
+    // If the import is already present in a top-level block, don't add it.
+    if python_ast
+        .iter()
+        .any(|stmt| includes_import(stmt, required_import))
+    {
+        return None;
+    }
+
     // Always insert the diagnostic at top-of-file.
     let mut diagnostic = Diagnostic::new(
         MissingRequiredImport(required_import.to_string()),
@@ -132,7 +129,6 @@ fn add_required_import(
 
 /// I002
 pub(crate) fn add_required_imports(
-    blocks: &[&Block],
     python_ast: &Suite,
     locator: &Locator,
     stylist: &Stylist,
@@ -174,7 +170,6 @@ pub(crate) fn add_required_imports(
                                 },
                                 level: level.map(|level| level.to_u32()),
                             }),
-                            blocks,
                             python_ast,
                             locator,
                             stylist,
@@ -193,7 +188,6 @@ pub(crate) fn add_required_imports(
                                     as_name: name.asname.as_deref(),
                                 },
                             }),
-                            blocks,
                             python_ast,
                             locator,
                             stylist,

crates/ruff/src/rules/isort/rules/organize_imports.rs

@@ -16,7 +16,7 @@ use ruff_python_ast::whitespace::leading_space;
 use crate::registry::AsRule;
 use crate::settings::Settings;
 
-use super::super::track::Block;
+use super::super::block::Block;
 use super::super::{comments, format_imports};
 
 /// ## What it does

crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__required_import_comment.py.snap

@@ -0,0 +1,19 @@
+---
+source: crates/ruff/src/rules/isort/mod.rs
+---
+comment.py:1:1: I002 [*] Missing required import: `from __future__ import annotations`
+  |
+1 | #!/usr/bin/env python3
+  |  I002
+2 | 
+3 | x = 1
+  |
+  = help: Insert required import: `from future import annotations`
+
+ℹ Suggested fix
+1 1 | #!/usr/bin/env python3
+  2 |+from __future__ import annotations
+2 3 | 
+3 4 | x = 1
+
+

crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__required_import_docstring_with_continuation.py.snap

@@ -0,0 +1,17 @@
+---
+source: crates/ruff/src/rules/isort/mod.rs
+---
+docstring_with_continuation.py:1:1: I002 [*] Missing required import: `from __future__ import annotations`
+  |
+1 | """Hello, world!"""; x = \
+  |  I002
+2 |     1; y = 2
+  |
+  = help: Insert required import: `from future import annotations`
+
+ℹ Suggested fix
+1   |-"""Hello, world!"""; x = \
+  1 |+"""Hello, world!"""; from __future__ import annotations; x = \
+2 2 |     1; y = 2
+
+

crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__required_import_docstring_with_semicolon.py.snap

@@ -0,0 +1,15 @@
+---
+source: crates/ruff/src/rules/isort/mod.rs
+---
+docstring_with_semicolon.py:1:1: I002 [*] Missing required import: `from __future__ import annotations`
+  |
+1 | """Hello, world!"""; x = 1
+  |  I002
+  |
+  = help: Insert required import: `from future import annotations`
+
+ℹ Suggested fix
+1   |-"""Hello, world!"""; x = 1
+  1 |+"""Hello, world!"""; from __future__ import annotations; x = 1
+
+

crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__required_import_existing_import.py.snap

@@ -0,0 +1,17 @@
+---
+source: crates/ruff/src/rules/isort/mod.rs
+---
+existing_import.py:1:1: I002 [*] Missing required import: `from __future__ import annotations`
+  |
+1 | from __future__ import generator_stop
+  |  I002
+2 | import os
+  |
+  = help: Insert required import: `from future import annotations`
+
+ℹ Suggested fix
+  1 |+from __future__ import annotations
+1 2 | from __future__ import generator_stop
+2 3 | import os
+
+

crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__required_import_off.py.snap

@@ -0,0 +1,20 @@
+---
+source: crates/ruff/src/rules/isort/mod.rs
+---
+off.py:1:1: I002 [*] Missing required import: `from __future__ import annotations`
+  |
+1 | # isort: off
+  |  I002
+2 | 
+3 | x = 1
+  |
+  = help: Insert required import: `from future import annotations`
+
+ℹ Suggested fix
+1 1 | # isort: off
+  2 |+from __future__ import annotations
+2 3 | 
+3 4 | x = 1
+4 5 | # isort: on
+
+