From 38fa305f358df8420f94ce62d072908010a38ec4 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sat, 8 Jul 2023 15:05:44 -0400 Subject: [PATCH] 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. --- .../resources/test/fixtures/isort/skip.py | 6 ++ .../resources/test/fixtures/isort/split.py | 10 +++ crates/ruff/src/rules/isort/block.rs | 76 +++++++++++-------- .../ruff__rules__isort__tests__skip.py.snap | 23 ++++++ .../ruff__rules__isort__tests__split.py.snap | 24 ++++++ 5 files changed, 109 insertions(+), 30 deletions(-) diff --git a/crates/ruff/resources/test/fixtures/isort/skip.py b/crates/ruff/resources/test/fixtures/isort/skip.py index 5cf668879bc99..97d38fdb155ff 100644 --- a/crates/ruff/resources/test/fixtures/isort/skip.py +++ b/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 diff --git a/crates/ruff/resources/test/fixtures/isort/split.py b/crates/ruff/resources/test/fixtures/isort/split.py index e4beaec56334a..c82885853a3a2 100644 --- a/crates/ruff/resources/test/fixtures/isort/split.py +++ b/crates/ruff/resources/test/fixtures/isort/split.py @@ -19,3 +19,13 @@ import D import B + + +import e +import f + +# isort: split +# isort: split + +import d +import c diff --git a/crates/ruff/src/rules/isort/block.rs b/crates/ruff/src/rules/isort/block.rs index 73ee365186524..e0eea8ae35eae 100644 --- a/crates/ruff/src/rules/isort/block.rs +++ b/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>, - splits: &'a [TextSize], - cell_offsets: Option<&'a [TextSize]>, + splits: Peekable>, + cell_offsets: Option>>, 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. - for (index, split) in self.splits.iter().enumerate() { - if stmt.start() >= *split { - self.finalize(self.trailer_for(stmt)); - self.splits = &self.splits[index + 1..]; - } else { - break; + // Track manual splits (e.g., `# isort: split`). + if self + .splits + .next_if(|split| stmt.start() >= **split) + .is_some() + { + // 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 { - for (index, split) in cell_offsets.iter().enumerate() { - if stmt.start() >= *split { - // We don't want any extra newlines between cells. - self.finalize(None); - self.cell_offsets = Some(&cell_offsets[index + 1..]); - } else { - break; + if let Some(cell_offsets) = self.cell_offsets.as_mut() { + if cell_offsets + .next_if(|cell_offset| stmt.start() >= **cell_offset) + .is_some() + { + // Skip any other cell offsets that occur before the current statement (e.g., in + // the case of multiple empty cells). + while cell_offsets + .peek() + .map_or(false, |split| stmt.start() >= **split) + { + cell_offsets.next(); } - } - } - // Test if the statement is in an excluded range - 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; + self.finalize(None); } } // 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)); diff --git a/crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__skip.py.snap b/crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__skip.py.snap index d3eb858cc5c07..498421b65bb8d 100644 --- a/crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__skip.py.snap +++ b/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 diff --git a/crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__split.py.snap b/crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__split.py.snap index 185696f70c41b..7556899b60c05 100644 --- a/crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__split.py.snap +++ b/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