Refactor isort directive skips to use iterators (#5623)

## Summary

We're doing some unsafe accesses to advance these iterators. It's easier
to model these as actual iterators to ensure safety everywhere. Also
added some additional test cases.

Closes #5621.

crates/ruff/resources/test/fixtures/isort/skip.py

@@ -26,3 +26,9 @@ def f():
     import os  # isort:skip
     import collections
     import abc
+
+
+def f():
+    import sys; import os  # isort:skip
+    import sys; import os  # isort:skip  # isort:skip
+    import sys; import os

crates/ruff/resources/test/fixtures/isort/split.py

@@ -19,3 +19,13 @@ if True:
 
     import D
     import B
+
+
+import e
+import f
+
+# isort: split
+# isort: split
+
+import d
+import c

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

@@ -1,5 +1,7 @@
 use ruff_text_size::{TextRange, TextSize};
 use rustpython_parser::ast::{self, ExceptHandler, MatchCase, Ranged, Stmt};
+use std::iter::Peekable;
+use std::slice;
 
 use ruff_python_ast::source_code::Locator;
 use ruff_python_ast::statement_visitor::StatementVisitor;
@@ -30,8 +32,8 @@ pub(crate) struct BlockBuilder<'a> {
     locator: &'a Locator<'a>,
     is_stub: bool,
     blocks: Vec<Block<'a>>,
-    splits: &'a [TextSize],
+    splits: Peekable<slice::Iter<'a, TextSize>>,
-    cell_offsets: Option<&'a [TextSize]>,
+    cell_offsets: Option<Peekable<slice::Iter<'a, TextSize>>>,
     exclusions: &'a [TextRange],
     nested: bool,
 }
@@ -47,12 +49,13 @@ impl<'a> BlockBuilder<'a> {
             locator,
             is_stub,
             blocks: vec![Block::default()],
-            splits: &directives.splits,
+            splits: directives.splits.iter().peekable(),
             exclusions: &directives.exclusions,
             nested: false,
             cell_offsets: source_kind
                 .and_then(SourceKind::notebook)
-                .map(Notebook::cell_offsets),
+                .map(Notebook::cell_offsets)
+                .map(|offsets| offsets.iter().peekable()),
         }
     }
 
@@ -126,45 +129,58 @@ where
     'b: 'a,
 {
     fn visit_stmt(&mut self, stmt: &'b Stmt) {
-        // Track manual splits.
+        // Track manual splits (e.g., `# isort: split`).
-        for (index, split) in self.splits.iter().enumerate() {
+        if self
-            if stmt.start() >= *split {
+            .splits
-                self.finalize(self.trailer_for(stmt));
+            .next_if(|split| stmt.start() >= **split)
-                self.splits = &self.splits[index + 1..];
+            .is_some()
-            } else {
+        {
-                break;
+            // Skip any other splits that occur before the current statement, to support, e.g.:
+            // ```python
+            // # isort: split
+            // # isort: split
+            // import foo
+            // ```
+            while self
+                .splits
+                .peek()
+                .map_or(false, |split| stmt.start() >= **split)
+            {
+                self.splits.next();
             }
+
+            self.finalize(self.trailer_for(stmt));
         }
 
         // Track Jupyter notebook cell offsets as splits. This will make sure
         // that each cell is considered as an individual block to organize the
         // imports in. Thus, not creating an edit which spans across multiple
         // cells.
-        if let Some(cell_offsets) = self.cell_offsets {
+        if let Some(cell_offsets) = self.cell_offsets.as_mut() {
-            for (index, split) in cell_offsets.iter().enumerate() {
+            if cell_offsets
-                if stmt.start() >= *split {
+                .next_if(|cell_offset| stmt.start() >= **cell_offset)
-                    // We don't want any extra newlines between cells.
+                .is_some()
-                    self.finalize(None);
+            {
-                    self.cell_offsets = Some(&cell_offsets[index + 1..]);
+                // Skip any other cell offsets that occur before the current statement (e.g., in
-                } else {
+                // the case of multiple empty cells).
-                    break;
+                while cell_offsets
+                    .peek()
+                    .map_or(false, |split| stmt.start() >= **split)
+                {
+                    cell_offsets.next();
                 }
-            }
-        }
 
-        // Test if the statement is in an excluded range
+                self.finalize(None);
-        let mut is_excluded = false;
-        for (index, exclusion) in self.exclusions.iter().enumerate() {
-            if exclusion.end() < stmt.start() {
-                self.exclusions = &self.exclusions[index + 1..];
-            } else {
-                is_excluded = exclusion.contains(stmt.start());
-                break;
             }
         }
 
         // Track imports.
-        if matches!(stmt, Stmt::Import(_) | Stmt::ImportFrom(_)) && !is_excluded {
+        if matches!(stmt, Stmt::Import(_) | Stmt::ImportFrom(_))
+            && !self
+                .exclusions
+                .iter()
+                .any(|exclusion| exclusion.contains(stmt.start()))
+        {
             self.track_import(stmt);
         } else {
             self.finalize(self.trailer_for(stmt));

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

@@ -31,6 +31,10 @@ skip.py:27:1: I001 [*] Import block is un-sorted or un-formatted
 26 |       import os  # isort:skip
 27 | /     import collections
 28 | |     import abc
+29 | | 
+   | |_^ I001
+30 |   
+31 |   def f():
    |
    = help: Organize imports
 
@@ -41,5 +45,24 @@ skip.py:27:1: I001 [*] Import block is un-sorted or un-formatted
    27 |+    import abc
 27 28 |     import collections
 28    |-    import abc
+29 29 | 
+30 30 | 
+31 31 | def f():
+
+skip.py:34:1: I001 [*] Import block is un-sorted or un-formatted
+   |
+32 |       import sys; import os  # isort:skip
+33 |       import sys; import os  # isort:skip  # isort:skip
+34 | /     import sys; import os
+   |
+   = help: Organize imports
+
+ℹ Fix
+31 31 | def f():
+32 32 |     import sys; import os  # isort:skip
+33 33 |     import sys; import os  # isort:skip  # isort:skip
+34    |-    import sys; import os
+   34 |+    import os
+   35 |+    import sys
 
 

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

@@ -29,6 +29,10 @@ split.py:20:1: I001 [*] Import block is un-sorted or un-formatted
 19 |   
 20 | /     import D
 21 | |     import B
+22 | | 
+   | |_^ I001
+23 |   
+24 |   import e
    |
    = help: Organize imports
 
@@ -39,5 +43,25 @@ split.py:20:1: I001 [*] Import block is un-sorted or un-formatted
    20 |+    import B
 20 21 |     import D
 21    |-    import B
+22 22 | 
+23 23 | 
+24 24 | import e
+
+split.py:30:1: I001 [*] Import block is un-sorted or un-formatted
+   |
+28 |   # isort: split
+29 |   
+30 | / import d
+31 | | import c
+   |
+   = help: Organize imports
+
+ℹ Fix
+27 27 | # isort: split
+28 28 | # isort: split
+29 29 | 
+   30 |+import c
+30 31 | import d
+31    |-import c