Import compatibility with isort newline-insertion behavior (#1078)

resources/test/fixtures/isort/insert_empty_lines.py

@@ -0,0 +1,16 @@
+import a
+import b
+x = 1
+import os
+import sys
+def f():
+    pass
+if True:
+    x = 1
+    import collections
+    import typing
+class X: pass
+y = 1
+import os
+import sys
+"""Docstring"""

src/ast/helpers.rs

@@ -301,6 +301,16 @@ pub fn match_trailing_content(stmt: &Stmt, locator: &SourceCodeLocator) -> bool
     false
 }
 
+/// Return the number of trailing empty lines following a statement.
+pub fn count_trailing_lines(stmt: &Stmt, locator: &SourceCodeLocator) -> usize {
+    let suffix =
+        locator.slice_source_code_at(&Location::new(stmt.end_location.unwrap().row() + 1, 0));
+    suffix
+        .lines()
+        .take_while(|line| line.trim().is_empty())
+        .count()
+}
+
 #[cfg(test)]
 mod tests {
     use anyhow::Result;

src/autofix/fixer.rs

@@ -92,7 +92,7 @@ fn apply_fixes<'a>(
     }
 
     // Add the remaining content.
-    let slice = locator.slice_source_code_at(last_pos);
+    let slice = locator.slice_source_code_at(&last_pos);
     output.append(&slice);
 
     (Cow::from(output.finish()), num_fixed)

src/check_imports.rs

@@ -18,7 +18,7 @@ fn check_import_blocks(
 ) -> Vec<Check> {
     let mut checks = vec![];
     for block in tracker.into_iter() {
-        if !block.is_empty() {
+        if !block.imports.is_empty() {
             if let Some(check) = isort::plugins::check_imports(&block, locator, settings, autofix) {
                 checks.push(check);
             }

src/isort/mod.rs

@@ -10,6 +10,7 @@ use rustpython_ast::{Stmt, StmtKind};
 use crate::isort::categorize::{categorize, ImportType};
 use crate::isort::comments::Comment;
 use crate::isort::sorting::{member_key, module_key};
+use crate::isort::track::{Block, Trailer};
 use crate::isort::types::{
     AliasData, CommentSet, ImportBlock, ImportFromData, Importable, OrderedImportBlock,
 };
@@ -464,7 +465,7 @@ fn sort_imports(block: ImportBlock) -> OrderedImportBlock {
 
 #[allow(clippy::too_many_arguments)]
 pub fn format_imports(
-    block: &[&Stmt],
+    block: &Block,
     comments: Vec<Comment>,
     line_length: usize,
     src: &[PathBuf],
@@ -474,7 +475,8 @@ pub fn format_imports(
     combine_as_imports: bool,
     force_wrap_aliases: bool,
 ) -> String {
-    let block = annotate_imports(block, comments);
+    let trailer = &block.trailer;
+    let block = annotate_imports(&block.imports, comments);
 
     // Normalize imports (i.e., deduplicate, aggregate `from` imports).
     let block = normalize_imports(block, combine_as_imports);
@@ -523,6 +525,16 @@ pub fn format_imports(
             is_first_statement = false;
         }
     }
+    match trailer {
+        None => {}
+        Some(Trailer::Sibling) => {
+            output.append("\n");
+        }
+        Some(Trailer::FunctionDef | Trailer::ClassDef) => {
+            output.append("\n");
+            output.append("\n");
+        }
+    }
     output.finish().to_string()
 }
 
@@ -546,6 +558,7 @@ mod tests {
     #[test_case(Path::new("fit_line_length_comment.py"))]
     #[test_case(Path::new("force_wrap_aliases.py"))]
     #[test_case(Path::new("import_from_after_import.py"))]
+    #[test_case(Path::new("insert_empty_lines.py"))]
     #[test_case(Path::new("leading_prefix.py"))]
     #[test_case(Path::new("no_reorder_within_section.py"))]
     #[test_case(Path::new("order_by_type.py"))]

src/isort/plugins.rs

@@ -1,11 +1,12 @@
 use rustpython_ast::{Location, Stmt};
 use textwrap::{dedent, indent};
 
-use crate::ast::helpers::{match_leading_content, match_trailing_content};
+use crate::ast::helpers::{count_trailing_lines, match_leading_content, match_trailing_content};
 use crate::ast::types::Range;
 use crate::ast::whitespace::leading_space;
 use crate::autofix::Fix;
 use crate::checks::CheckKind;
+use crate::isort::track::Block;
 use crate::isort::{comments, format_imports};
 use crate::{Check, Settings, SourceCodeLocator};
 
@@ -30,13 +31,13 @@ fn extract_indentation(body: &[&Stmt], locator: &SourceCodeLocator) -> String {
 
 /// I001
 pub fn check_imports(
-    body: &[&Stmt],
+    block: &Block,
     locator: &SourceCodeLocator,
     settings: &Settings,
     autofix: bool,
 ) -> Option<Check> {
-    let range = extract_range(body);
-    let indentation = extract_indentation(body, locator);
+    let range = extract_range(&block.imports);
+    let indentation = extract_indentation(&block.imports, locator);
 
     // Extract comments. Take care to grab any inline comments from the last line.
     let comments = comments::collect_comments(
@@ -48,12 +49,13 @@ pub fn check_imports(
     );
 
     // Special-cases: there's leading or trailing content in the import block.
-    let has_leading_content = match_leading_content(body.first().unwrap(), locator);
-    let has_trailing_content = match_trailing_content(body.last().unwrap(), locator);
+    let has_leading_content = match_leading_content(block.imports.first().unwrap(), locator);
+    let has_trailing_content = match_trailing_content(block.imports.last().unwrap(), locator);
+    let num_trailing_lines = count_trailing_lines(block.imports.last().unwrap(), locator);
 
     // Generate the sorted import block.
     let expected = format_imports(
-        body,
+        block,
         comments,
         settings.line_length - indentation.len(),
         &settings.src,
@@ -81,7 +83,7 @@ pub fn check_imports(
                     Location::new(range.location.row(), 0)
                 },
                 // TODO(charlie): Preserve trailing suffixes. Right now, we strip them.
-                Location::new(range.end_location.row() + 1, 0),
+                Location::new(range.end_location.row() + 1 + num_trailing_lines, 0),
             ));
         }
         Some(check)
@@ -89,7 +91,7 @@ pub fn check_imports(
         // Expand the span the entire range, including leading and trailing space.
         let range = Range {
             location: Location::new(range.location.row(), 0),
-            end_location: Location::new(range.end_location.row() + 1, 0),
+            end_location: Location::new(range.end_location.row() + 1 + num_trailing_lines, 0),
         };
         let actual = dedent(&locator.slice_source_code_range(&range));
         if actual == expected {

src/isort/snapshots/ruff__isort__tests__insert_empty_lines.py.snap

@@ -0,0 +1,50 @@
+---
+source: src/isort/mod.rs
+expression: checks
+---
+- kind: UnsortedImports
+  location:
+    row: 1
+    column: 0
+  end_location:
+    row: 3
+    column: 0
+  fix:
+    content: "import a\nimport b\n\n"
+    location:
+      row: 1
+      column: 0
+    end_location:
+      row: 3
+      column: 0
+- kind: UnsortedImports
+  location:
+    row: 4
+    column: 0
+  end_location:
+    row: 6
+    column: 0
+  fix:
+    content: "import os\nimport sys\n\n\n"
+    location:
+      row: 4
+      column: 0
+    end_location:
+      row: 6
+      column: 0
+- kind: UnsortedImports
+  location:
+    row: 14
+    column: 0
+  end_location:
+    row: 16
+    column: 0
+  fix:
+    content: "import os\nimport sys\n\n"
+    location:
+      row: 14
+      column: 0
+    end_location:
+      row: 16
+      column: 0
+

src/isort/snapshots/ruff__isort__tests__leading_prefix.py.snap

@@ -10,12 +10,12 @@ expression: checks
     row: 2
     column: 9
   fix:
-    content: "\nimport os\nimport sys\n"
+    content: "\nimport os\nimport sys\n\n"
     location:
       row: 1
       column: 7
     end_location:
-      row: 3
+      row: 4
       column: 0
 - kind: UnsortedImports
   location:

src/isort/snapshots/ruff__isort__tests__skip.py.snap

@@ -2,6 +2,21 @@
 source: src/isort/mod.rs
 expression: checks
 ---
+- kind: UnsortedImports
+  location:
+    row: 7
+    column: 0
+  end_location:
+    row: 8
+    column: 0
+  fix:
+    content: "import sys\n\n"
+    location:
+      row: 7
+      column: 0
+    end_location:
+      row: 8
+      column: 0
 - kind: UnsortedImports
   location:
     row: 9

src/isort/snapshots/ruff__isort__tests__trailing_suffix.py.snap

@@ -10,12 +10,12 @@ expression: checks
     row: 2
     column: 9
   fix:
-    content: "import os\nimport sys\n"
+    content: "import os\nimport sys\n\n"
     location:
       row: 1
       column: 0
     end_location:
-      row: 3
+      row: 4
       column: 0
 - kind: UnsortedImports
   location:
@@ -25,7 +25,7 @@ expression: checks
     row: 6
     column: 13
   fix:
-    content: "    import os\n    import sys\n"
+    content: "    import os\n    import sys\n\n"
     location:
       row: 5
       column: 0

src/isort/track.rs

@@ -7,33 +7,47 @@ use rustpython_ast::{
 
 use crate::ast::visitor::Visitor;
 
+#[derive(Debug)]
+pub enum Trailer {
+    Sibling,
+    ClassDef,
+    FunctionDef,
+}
+
+#[derive(Debug, 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<Vec<&'a Stmt>>,
+    blocks: Vec<Block<'a>>,
 }
 
 impl<'a> ImportTracker<'a> {
     pub fn new(exclusions: &'a IntSet<usize>) -> Self {
         Self {
             exclusions,
-            blocks: vec![vec![]],
+            blocks: vec![Block::default()],
         }
     }
 
     fn track_import(&mut self, stmt: &'a Stmt) {
         let index = self.blocks.len() - 1;
-        self.blocks[index].push(stmt);
+        self.blocks[index].imports.push(stmt);
     }
 
-    fn finalize(&mut self) {
+    fn finalize(&mut self, trailer: Option<Trailer>) {
         let index = self.blocks.len() - 1;
-        if !self.blocks[index].is_empty() {
-            self.blocks.push(vec![]);
+        if !self.blocks[index].imports.is_empty() {
+            self.blocks[index].trailer = trailer;
+            self.blocks.push(Block::default());
         }
     }
 
-    pub fn into_iter(self) -> impl IntoIterator<Item = Vec<&'a Stmt>> {
+    pub fn into_iter(self) -> impl IntoIterator<Item = Block<'a>> {
         self.blocks.into_iter()
     }
 }
@@ -51,7 +65,13 @@ where
         {
             self.track_import(stmt);
         } else {
-            self.finalize();
+            self.finalize(Some(match &stmt.node {
+                StmtKind::FunctionDef { .. } | StmtKind::AsyncFunctionDef { .. } => {
+                    Trailer::FunctionDef
+                }
+                StmtKind::ClassDef { .. } => Trailer::ClassDef,
+                _ => Trailer::Sibling,
+            }));
         }
 
         // Track scope.
@@ -60,75 +80,75 @@ where
                 for stmt in body {
                     self.visit_stmt(stmt);
                 }
-                self.finalize();
+                self.finalize(None);
             }
             StmtKind::AsyncFunctionDef { body, .. } => {
                 for stmt in body {
                     self.visit_stmt(stmt);
                 }
-                self.finalize();
+                self.finalize(None);
             }
             StmtKind::ClassDef { body, .. } => {
                 for stmt in body {
                     self.visit_stmt(stmt);
                 }
-                self.finalize();
+                self.finalize(None);
             }
             StmtKind::For { body, orelse, .. } => {
                 for stmt in body {
                     self.visit_stmt(stmt);
                 }
-                self.finalize();
+                self.finalize(None);
 
                 for stmt in orelse {
                     self.visit_stmt(stmt);
                 }
-                self.finalize();
+                self.finalize(None);
             }
             StmtKind::AsyncFor { body, orelse, .. } => {
                 for stmt in body {
                     self.visit_stmt(stmt);
                 }
-                self.finalize();
+                self.finalize(None);
 
                 for stmt in orelse {
                     self.visit_stmt(stmt);
                 }
-                self.finalize();
+                self.finalize(None);
             }
             StmtKind::While { body, orelse, .. } => {
                 for stmt in body {
                     self.visit_stmt(stmt);
                 }
-                self.finalize();
+                self.finalize(None);
 
                 for stmt in orelse {
                     self.visit_stmt(stmt);
                 }
-                self.finalize();
+                self.finalize(None);
             }
             StmtKind::If { body, orelse, .. } => {
                 for stmt in body {
                     self.visit_stmt(stmt);
                 }
-                self.finalize();
+                self.finalize(None);
 
                 for stmt in orelse {
                     self.visit_stmt(stmt);
                 }
-                self.finalize();
+                self.finalize(None);
             }
             StmtKind::With { body, .. } => {
                 for stmt in body {
                     self.visit_stmt(stmt);
                 }
-                self.finalize();
+                self.finalize(None);
             }
             StmtKind::AsyncWith { body, .. } => {
                 for stmt in body {
                     self.visit_stmt(stmt);
                 }
-                self.finalize();
+                self.finalize(None);
             }
             StmtKind::Match { cases, .. } => {
                 for match_case in cases {
@@ -148,17 +168,17 @@ where
                 for stmt in body {
                     self.visit_stmt(stmt);
                 }
-                self.finalize();
+                self.finalize(None);
 
                 for stmt in orelse {
                     self.visit_stmt(stmt);
                 }
-                self.finalize();
+                self.finalize(None);
 
                 for stmt in finalbody {
                     self.visit_stmt(stmt);
                 }
-                self.finalize();
+                self.finalize(None);
             }
             _ => {}
         }
@@ -187,7 +207,7 @@ where
         for stmt in body {
             self.visit_stmt(stmt);
         }
-        self.finalize();
+        self.finalize(None);
     }
 
     fn visit_arguments(&mut self, _: &'b Arguments) {}
@@ -204,7 +224,7 @@ where
         for stmt in &match_case.body {
             self.visit_stmt(stmt);
         }
-        self.finalize();
+        self.finalize(None);
     }
 
     fn visit_pattern(&mut self, _: &'b Pattern) {}

src/pyupgrade/fixes.rs

@@ -21,7 +21,7 @@ pub fn remove_class_def_base(
     bases: &[Expr],
     keywords: &[Keyword],
 ) -> Option<Fix> {
-    let contents = locator.slice_source_code_at(stmt_at);
+    let contents = locator.slice_source_code_at(&stmt_at);
 
     // Case 1: `object` is the only base.
     if bases.len() == 1 && keywords.is_empty() {

src/source_code_locator.rs

@@ -25,7 +25,7 @@ impl<'a> SourceCodeLocator<'a> {
         self.rope.get_or_init(|| Rope::from_str(self.contents))
     }
 
-    pub fn slice_source_code_at(&self, location: Location) -> Cow<'_, str> {
+    pub fn slice_source_code_at(&self, location: &Location) -> Cow<'_, str> {
         let rope = self.get_or_init_rope();
         let offset = rope.line_to_char(location.row() - 1) + location.column();
         Cow::from(rope.slice(offset..))