Avoid inserting extra newlines for comment-delimited import blocks (#1201)

resources/test/fixtures/isort/insert_empty_lines.py

@@ -39,3 +39,27 @@ if True:
     import collections
     import typing
     def f(): pass
+
+
+import os
+
+# Comment goes here.
+
+def f():
+    pass
+
+
+import os
+
+# Comment goes here.
+def f():
+    pass
+
+
+import os
+
+# Comment goes here.
+
+# And another.
+def f():
+    pass

resources/test/fixtures/isort/insert_empty_lines.pyi

@@ -39,3 +39,27 @@ if True:
     import collections
     import typing
     def f(): pass
+
+
+import os
+
+# Comment goes here.
+
+def f():
+    pass
+
+
+import os
+
+# Comment goes here.
+def f():
+    pass
+
+
+import os
+
+# Comment goes here.
+
+# And another.
+def f():
+    pass

src/check_imports.rs

@@ -37,7 +37,7 @@ pub fn check_imports(
     autofix: bool,
     path: &Path,
 ) -> Vec<Check> {
-    let mut tracker = ImportTracker::new(directives, path);
+    let mut tracker = ImportTracker::new(locator, directives, path);
     for stmt in python_ast {
         tracker.visit_stmt(stmt);
     }

src/isort/helpers.rs

@@ -0,0 +1,55 @@
+use rustpython_ast::Stmt;
+
+use crate::source_code_locator::SourceCodeLocator;
+
+/// Return `true` if a `Stmt` is preceded by a "comment break"
+pub fn has_comment_break(stmt: &Stmt, locator: &SourceCodeLocator) -> bool {
+    // Starting from the `Stmt` (`def f(): pass`), we want to detect patterns like
+    // this:
+    //
+    //   import os
+    //
+    //   # Detached comment.
+    //
+    //   def f(): pass
+
+    // This should also be detected:
+    //
+    //   import os
+    //
+    //   # Detached comment.
+    //
+    //   # Direct comment.
+    //   def f(): pass
+
+    // But this should not:
+    //
+    //   import os
+    //
+    //   # Direct comment.
+    //   def f(): pass
+    let mut seen_blank = false;
+    for line in locator
+        .slice_source_code_until(&stmt.location)
+        .lines()
+        .rev()
+    {
+        let line = line.trim();
+        if seen_blank {
+            if line.starts_with('#') {
+                return true;
+            } else if !line.is_empty() {
+                break;
+            }
+        } else {
+            if line.is_empty() {
+                seen_blank = true;
+            } else if line.starts_with('#') || line.starts_with('@') {
+                continue;
+            } else {
+                break;
+            }
+        }
+    }
+    false
+}

src/isort/mod.rs

@@ -18,6 +18,7 @@ use crate::isort::types::{
 mod categorize;
 mod comments;
 pub mod format;
+mod helpers;
 pub mod plugins;
 pub mod settings;
 mod sorting;

src/isort/snapshots/ruff__isort__tests__insert_empty_lines.py.snap

@@ -49,32 +49,17 @@ expression: checks
       column: 0
 - kind: UnsortedImports
   location:
-    row: 33
+    row: 52
     column: 0
   end_location:
-    row: 35
+    row: 54
     column: 0
   fix:
-    content: "    import collections\n    import typing\n\n"
+    content: "import os\n\n\n"
     location:
-      row: 33
+      row: 52
       column: 0
     end_location:
-      row: 35
-      column: 0
-- kind: UnsortedImports
-  location:
-    row: 39
-    column: 0
-  end_location:
-    row: 41
-    column: 0
-  fix:
-    content: "    import collections\n    import typing\n\n"
-    location:
-      row: 39
-      column: 0
-    end_location:
-      row: 41
+      row: 54
       column: 0
 

src/isort/snapshots/ruff__isort__tests__insert_empty_lines.pyi.snap

@@ -47,19 +47,4 @@ expression: checks
     end_location:
       row: 16
       column: 0
-- kind: UnsortedImports
-  location:
-    row: 33
-    column: 0
-  end_location:
-    row: 35
-    column: 0
-  fix:
-    content: "    import collections\n    import typing\n\n"
-    location:
-      row: 33
-      column: 0
-    end_location:
-      row: 35
-      column: 0
 

src/isort/snapshots/ruff__isort__tests__trailing_suffix.py.snap

@@ -25,7 +25,7 @@ expression: checks
     row: 6
     column: 13
   fix:
-    content: "    import os\n    import sys\n\n"
+    content: "    import os\n    import sys\n"
     location:
       row: 5
       column: 0

src/isort/track.rs

@@ -8,20 +8,24 @@ use rustpython_ast::{
 
 use crate::ast::visitor::Visitor;
 use crate::directives::IsortDirectives;
+use crate::isort::helpers;
+use crate::source_code_locator::SourceCodeLocator;
 
+#[derive(Debug)]
 pub enum Trailer {
     Sibling,
     ClassDef,
     FunctionDef,
 }
 
-#[derive(Default)]
+#[derive(Debug, Default)]
 pub struct Block<'a> {
     pub imports: Vec<&'a Stmt>,
     pub trailer: Option<Trailer>,
 }
 
 pub struct ImportTracker<'a> {
+    locator: &'a SourceCodeLocator<'a>,
     directives: &'a IsortDirectives,
     pyi: bool,
     blocks: Vec<Block<'a>>,
@@ -30,8 +34,13 @@ pub struct ImportTracker<'a> {
 }
 
 impl<'a> ImportTracker<'a> {
-    pub fn new(directives: &'a IsortDirectives, path: &'a Path) -> Self {
+    pub fn new(
+        locator: &'a SourceCodeLocator<'a>,
+        directives: &'a IsortDirectives,
+        path: &'a Path,
+    ) -> Self {
         Self {
+            locator,
             directives,
             pyi: path.extension().map_or(false, |ext| ext == "pyi"),
             blocks: vec![Block::default()],
@@ -46,31 +55,46 @@ impl<'a> ImportTracker<'a> {
     }
 
     fn trailer_for(&self, stmt: &'a Stmt) -> Option<Trailer> {
-        if self.pyi {
-            // Black treats interface files differently, limiting to one newline
-            // (`Trailing::Sibling`), and avoiding inserting any newlines in nested function
-            // blocks.
-            if self.nested
-                && matches!(
-                    stmt.node,
-                    StmtKind::FunctionDef { .. } | StmtKind::AsyncFunctionDef { .. }
-                )
-            {
-                None
-            } else {
-                Some(Trailer::Sibling)
-            }
-        } else if self.nested {
-            Some(Trailer::Sibling)
-        } else {
-            Some(match &stmt.node {
-                StmtKind::FunctionDef { .. } | StmtKind::AsyncFunctionDef { .. } => {
-                    Trailer::FunctionDef
-                }
-                StmtKind::ClassDef { .. } => Trailer::ClassDef,
-                _ => Trailer::Sibling,
-            })
+        // No need to compute trailers if we won't be finalizing anything.
+        let index = self.blocks.len() - 1;
+        if self.blocks[index].imports.is_empty() {
+            return None;
         }
+
+        // Similar to isort, avoid enforcing any newline behaviors in nested blocks.
+        if self.nested {
+            return None;
+        }
+
+        Some(if self.pyi {
+            // Black treats interface files differently, limiting to one newline
+            // (`Trailing::Sibling`).
+            Trailer::Sibling
+        } else {
+            // If the import block is followed by a class or function, we want to enforce
+            // two blank lines. The exception: if, between the import and the class or
+            // function, we have at least one commented line, followed by at
+            // least one blank line. In that case, we treat it as a regular
+            // sibling (i.e., as if the comment is the next statement, as
+            // opposed to the class or function).
+            match &stmt.node {
+                StmtKind::FunctionDef { .. } | StmtKind::AsyncFunctionDef { .. } => {
+                    if helpers::has_comment_break(stmt, self.locator) {
+                        Trailer::Sibling
+                    } else {
+                        Trailer::FunctionDef
+                    }
+                }
+                StmtKind::ClassDef { .. } => {
+                    if helpers::has_comment_break(stmt, self.locator) {
+                        Trailer::Sibling
+                    } else {
+                        Trailer::ClassDef
+                    }
+                }
+                _ => Trailer::Sibling,
+            }
+        })
     }
 
     fn finalize(&mut self, trailer: Option<Trailer>) {

src/source_code_locator.rs

@@ -31,6 +31,12 @@ impl<'a> SourceCodeLocator<'a> {
         Cow::from(rope.slice(offset..))
     }
 
+    pub fn slice_source_code_until(&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))
+    }
+
     pub fn slice_source_code_range(&self, range: &Range) -> Cow<'_, str> {
         let rope = self.get_or_init_rope();
         let start = rope.line_to_char(range.location.row() - 1) + range.location.column();