Support isort: split directive (#1081)

src/check_imports.rs

@@ -1,10 +1,10 @@
 //! Lint rules based on import analysis.
 
-use nohash_hasher::IntSet;
 use rustpython_parser::ast::Suite;
 
 use crate::ast::visitor::Visitor;
 use crate::checks::Check;
+use crate::directives::IsortDirectives;
 use crate::isort;
 use crate::isort::track::ImportTracker;
 use crate::settings::Settings;
@@ -30,11 +30,11 @@ fn check_import_blocks(
 pub fn check_imports(
     python_ast: &Suite,
     locator: &SourceCodeLocator,
-    exclusions: &IntSet<usize>,
+    directives: &IsortDirectives,
     settings: &Settings,
     autofix: bool,
 ) -> Vec<Check> {
-    let mut tracker = ImportTracker::new(exclusions);
+    let mut tracker = ImportTracker::new(directives);
     for stmt in python_ast {
         tracker.visit_stmt(stmt);
     }

src/directives.rs

@@ -30,9 +30,15 @@ impl Flags {
     }
 }
 
+#[derive(Default)]
+pub struct IsortDirectives {
+    pub exclusions: IntSet<usize>,
+    pub splits: Vec<usize>,
+}
+
 pub struct Directives {
     pub noqa_line_for: IntMap<usize, usize>,
-    pub isort_exclusions: IntSet<usize>,
+    pub isort: IsortDirectives,
 }
 
 pub fn extract_directives(
@@ -46,10 +52,10 @@ pub fn extract_directives(
         } else {
             IntMap::default()
         },
-        isort_exclusions: if flags.contains(Flags::ISORT) {
-            extract_isort_exclusions(lxr, locator)
+        isort: if flags.contains(Flags::ISORT) {
+            extract_isort_directives(lxr, locator)
         } else {
-            IntSet::default()
+            IsortDirectives::default()
         },
     }
 }
@@ -73,8 +79,9 @@ pub fn extract_noqa_line_for(lxr: &[LexResult]) -> IntMap<usize, usize> {
 }
 
 /// Extract a set of lines over which to disable isort.
-pub fn extract_isort_exclusions(lxr: &[LexResult], locator: &SourceCodeLocator) -> IntSet<usize> {
+pub fn extract_isort_directives(lxr: &[LexResult], locator: &SourceCodeLocator) -> IsortDirectives {
     let mut exclusions: IntSet<usize> = IntSet::default();
+    let mut splits: Vec<usize> = Vec::default();
     let mut skip_file: bool = false;
     let mut off: Option<Location> = None;
     let mut last: Option<Location> = None;
@@ -92,7 +99,10 @@ pub fn extract_isort_exclusions(lxr: &[LexResult], locator: &SourceCodeLocator)
                 location: start,
                 end_location: end,
             });
-            if comment_text == "# isort: skip_file" {
+
+            if comment_text == "# isort: split" {
+                splits.push(start.row());
+            } else if comment_text == "# isort: skip_file" {
                 skip_file = true;
             } else if off.is_some() {
                 if comment_text == "# isort: on" {
@@ -127,7 +137,7 @@ pub fn extract_isort_exclusions(lxr: &[LexResult], locator: &SourceCodeLocator)
             }
         }
     }
-    exclusions
+    IsortDirectives { exclusions, splits }
 }
 
 #[cfg(test)]
@@ -136,7 +146,7 @@ mod tests {
     use rustpython_parser::lexer;
     use rustpython_parser::lexer::LexResult;
 
-    use crate::directives::{extract_isort_exclusions, extract_noqa_line_for};
+    use crate::directives::{extract_isort_directives, extract_noqa_line_for};
     use crate::SourceCodeLocator;
 
     #[test]
@@ -220,13 +230,16 @@ z = x + 1",
     }
 
     #[test]
-    fn isort_extraction() {
+    fn isort_exclusions() {
         let contents = "x = 1
 y = 2
 z = x + 1";
         let lxr: Vec<LexResult> = lexer::make_tokenizer(contents).collect();
         let locator = SourceCodeLocator::new(contents);
-        assert_eq!(extract_isort_exclusions(&lxr, &locator), IntSet::default());
+        assert_eq!(
+            extract_isort_directives(&lxr, &locator).exclusions,
+            IntSet::default()
+        );
 
         let contents = "# isort: off
 x = 1
@@ -236,7 +249,7 @@ z = x + 1";
         let lxr: Vec<LexResult> = lexer::make_tokenizer(contents).collect();
         let locator = SourceCodeLocator::new(contents);
         assert_eq!(
-            extract_isort_exclusions(&lxr, &locator),
+            extract_isort_directives(&lxr, &locator).exclusions,
             IntSet::from_iter([2, 3, 4])
         );
 
@@ -250,7 +263,7 @@ z = x + 1
         let lxr: Vec<LexResult> = lexer::make_tokenizer(contents).collect();
         let locator = SourceCodeLocator::new(contents);
         assert_eq!(
-            extract_isort_exclusions(&lxr, &locator),
+            extract_isort_directives(&lxr, &locator).exclusions,
             IntSet::from_iter([2, 3, 4, 5])
         );
 
@@ -261,7 +274,7 @@ z = x + 1";
         let lxr: Vec<LexResult> = lexer::make_tokenizer(contents).collect();
         let locator = SourceCodeLocator::new(contents);
         assert_eq!(
-            extract_isort_exclusions(&lxr, &locator),
+            extract_isort_directives(&lxr, &locator).exclusions,
             IntSet::from_iter([2, 3, 4])
         );
 
@@ -272,7 +285,7 @@ z = x + 1";
         let lxr: Vec<LexResult> = lexer::make_tokenizer(contents).collect();
         let locator = SourceCodeLocator::new(contents);
         assert_eq!(
-            extract_isort_exclusions(&lxr, &locator),
+            extract_isort_directives(&lxr, &locator).exclusions,
             IntSet::from_iter([1, 2, 3, 4])
         );
 
@@ -285,8 +298,36 @@ z = x + 1";
         let lxr: Vec<LexResult> = lexer::make_tokenizer(contents).collect();
         let locator = SourceCodeLocator::new(contents);
         assert_eq!(
-            extract_isort_exclusions(&lxr, &locator),
+            extract_isort_directives(&lxr, &locator).exclusions,
             IntSet::from_iter([1, 2, 3, 4, 5, 6])
         );
     }
+
+    #[test]
+    fn isort_splits() {
+        let contents = "x = 1
+y = 2
+z = x + 1";
+        let lxr: Vec<LexResult> = lexer::make_tokenizer(contents).collect();
+        let locator = SourceCodeLocator::new(contents);
+        assert_eq!(
+            extract_isort_directives(&lxr, &locator).splits,
+            Vec::<usize>::new()
+        );
+
+        let contents = "x = 1
+y = 2
+# isort: split
+z = x + 1";
+        let lxr: Vec<LexResult> = lexer::make_tokenizer(contents).collect();
+        let locator = SourceCodeLocator::new(contents);
+        assert_eq!(extract_isort_directives(&lxr, &locator).splits, vec![3]);
+
+        let contents = "x = 1
+y = 2  # isort: split
+z = x + 1";
+        let lxr: Vec<LexResult> = lexer::make_tokenizer(contents).collect();
+        let locator = SourceCodeLocator::new(contents);
+        assert_eq!(extract_isort_directives(&lxr, &locator).splits, vec![2]);
+    }
 }

src/isort/mod.rs

@@ -572,7 +572,9 @@ mod tests {
     #[test_case(Path::new("separate_local_folder_imports.py"))]
     #[test_case(Path::new("separate_third_party_imports.py"))]
     #[test_case(Path::new("skip.py"))]
+    #[test_case(Path::new("skip_file.py"))]
     #[test_case(Path::new("sort_similar_imports.py"))]
+    #[test_case(Path::new("split.py"))]
     #[test_case(Path::new("trailing_suffix.py"))]
     #[test_case(Path::new("type_comments.py"))]
     fn default(path: &Path) -> Result<()> {

src/isort/snapshots/ruff__isort__tests__skip_file.py.snap

@@ -0,0 +1,6 @@
+---
+source: src/isort/mod.rs
+expression: checks
+---
+[]
+

src/isort/snapshots/ruff__isort__tests__split.py.snap

@@ -0,0 +1,6 @@
+---
+source: src/isort/mod.rs
+expression: checks
+---
+[]
+

src/isort/track.rs

@@ -1,4 +1,3 @@
-use nohash_hasher::IntSet;
 use rustpython_ast::{
     Alias, Arg, Arguments, Boolop, Cmpop, Comprehension, Constant, Excepthandler,
     ExcepthandlerKind, Expr, ExprContext, Keyword, MatchCase, Operator, Pattern, Stmt, StmtKind,
@@ -6,31 +5,32 @@ use rustpython_ast::{
 };
 
 use crate::ast::visitor::Visitor;
+use crate::directives::IsortDirectives;
 
-#[derive(Debug)]
 pub enum Trailer {
     Sibling,
     ClassDef,
     FunctionDef,
 }
 
-#[derive(Debug, Default)]
+#[derive(Default)]
 pub struct Block<'a> {
     pub imports: Vec<&'a Stmt>,
     pub trailer: Option<Trailer>,
 }
 
-#[derive(Debug)]
 pub struct ImportTracker<'a> {
-    exclusions: &'a IntSet<usize>,
     blocks: Vec<Block<'a>>,
+    directives: &'a IsortDirectives,
+    split_index: usize,
 }
 
 impl<'a> ImportTracker<'a> {
-    pub fn new(exclusions: &'a IntSet<usize>) -> Self {
+    pub fn new(directives: &'a IsortDirectives) -> Self {
         Self {
-            exclusions,
+            directives,
             blocks: vec![Block::default()],
+            split_index: 0,
         }
     }
 
@@ -57,11 +57,27 @@ where
     'b: 'a,
 {
     fn visit_stmt(&mut self, stmt: &'b Stmt) {
+        // Track manual splits.
+        while self.split_index < self.directives.splits.len() {
+            if stmt.location.row() >= self.directives.splits[self.split_index] {
+                self.finalize(Some(match &stmt.node {
+                    StmtKind::FunctionDef { .. } | StmtKind::AsyncFunctionDef { .. } => {
+                        Trailer::FunctionDef
+                    }
+                    StmtKind::ClassDef { .. } => Trailer::ClassDef,
+                    _ => Trailer::Sibling,
+                }));
+                self.split_index += 1;
+            } else {
+                break;
+            }
+        }
+
         // Track imports.
         if matches!(
             stmt.node,
             StmtKind::Import { .. } | StmtKind::ImportFrom { .. }
-        ) && !self.exclusions.contains(&stmt.location.row())
+        ) && !self.directives.exclusions.contains(&stmt.location.row())
         {
             self.track_import(stmt);
         } else {

src/linter.rs

@@ -89,7 +89,7 @@ pub(crate) fn check_path(
                     checks.extend(check_imports(
                         &python_ast,
                         locator,
-                        &directives.isort_exclusions,
+                        &directives.isort,
                         settings,
                         autofix,
                     ));
@@ -193,7 +193,7 @@ pub fn add_noqa_to_path(path: &Path, settings: &Settings) -> Result<usize> {
         &locator,
         &Directives {
             noqa_line_for: IntMap::default(),
-            isort_exclusions: directives.isort_exclusions,
+            isort: directives.isort,
         },
         settings,
         false,