Fix bugs

resources/test/fixtures/isort/add_newline_before_comments.py

@@ -0,0 +1,7 @@
+import os
+# This is a comment in the same section, so we need to add one newline.
+import sys
+import numpy as np
+# This is a comment, but it starts a new section, so we don't need to add a newline
+# before it.
+import leading_prefix

resources/test/fixtures/isort/fit_line_length_comment.py

@@ -0,0 +1,4 @@
+import a
+# Don't take this comment into account when determining whether the next import can fit on one line.
+from b import c
+from d import e  # Do take this comment into account when determining whether the next import can fit on one line.

resources/test/fixtures/isort/preserve_comment_order.py

@@ -0,0 +1,11 @@
+import io
+# Old MacDonald had a farm,
+# EIEIO
+# And on his farm he had a cow,
+# EIEIO
+# With a moo-moo here and a moo-moo there
+# Here a moo, there a moo, everywhere moo-moo
+# Old MacDonald had a farm,
+# EIEIO
+from errno import EIO
+import abc

src/isort/comments.rs

@@ -38,6 +38,3 @@ pub fn collect_comments<'a>(range: &Range, locator: &'a SourceCodeLocator) -> Ve
         })
         .collect()
 }
-
-// We have to assign each comment to a node. Then when we collect, we merge
-// nodes, or something. Comments can either be above, or to the right of a node.

src/isort/format.rs

@@ -6,7 +6,15 @@ use crate::isort::types::{AliasData, CommentSet, ImportFromData, Importable};
 const INDENT: &str = "    ";
 
 /// Add a plain import statement to the `RopeBuilder`.
-pub fn format_import(output: &mut RopeBuilder, alias: &AliasData, comments: &CommentSet) {
+pub fn format_import(
+    output: &mut RopeBuilder,
+    alias: &AliasData,
+    comments: &CommentSet,
+    is_first: bool,
+) {
+    if !is_first && !comments.atop.is_empty() {
+        output.append("\n");
+    }
     for comment in &comments.atop {
         output.append(&format!("{}\n", comment));
     }
@@ -28,6 +36,7 @@ pub fn format_import_from(
     comments: &CommentSet,
     aliases: &[(AliasData, CommentSet)],
     line_length: &usize,
+    is_first: bool,
 ) {
     // We can only inline if: (1) none of the aliases have atop comments, and (3)
     // only the last alias (if any) has inline comments.
@@ -40,17 +49,15 @@ pub fn format_import_from(
             .skip(1)
             .all(|(_, CommentSet { inline, .. })| inline.is_empty())
     {
-        // STOPSHIP(charlie): This includes the length of the comments...
-        let (single_line, import_length) = format_single_line(import_from, comments, aliases);
-        // If the import fits on a single line (excluding the newline character at the
-        // end, which doesn't count towards the line length), return it.
+        let (single_line, import_length) =
+            format_single_line(import_from, comments, aliases, is_first);
         if import_length <= *line_length {
             output.append(&single_line);
             return;
         }
     }
 
-    output.append(&format_multi_line(import_from, comments, aliases));
+    output.append(&format_multi_line(import_from, comments, aliases, is_first));
 }
 
 /// Format an import-from statement in single-line format.
@@ -60,55 +67,68 @@ fn format_single_line(
     import_from: &ImportFromData,
     comments: &CommentSet,
     aliases: &[(AliasData, CommentSet)],
+    is_first: bool,
 ) -> (String, usize) {
     let mut output = String::new();
-    let mut import_length = 0;
+    let mut line_length = 0;
 
+    if !is_first && !comments.atop.is_empty() {
+        output.push('\n');
+    }
     for comment in &comments.atop {
         output.push_str(comment);
         output.push('\n');
     }
 
-    output.push_str(&format!("from {} import ", import_from.module_name()));
+    let content = format!("from {} import ", import_from.module_name());
+    output.push_str(&content);
+    line_length += content.len();
 
     for (index, (AliasData { name, asname }, comments)) in aliases.iter().enumerate() {
-        for comment in &comments.atop {
-            output.push_str(comment);
-            output.push('\n');
-        }
         if let Some(asname) = asname {
             output.push_str(name);
             output.push_str(" as ");
             output.push_str(asname);
+            line_length += name.len() + 4 + asname.len();
         } else {
             output.push_str(name);
+            line_length += name.len();
         }
         if index < aliases.len() - 1 {
             output.push_str(", ");
+            line_length += 2;
         }
 
         for comment in &comments.inline {
-            output.push_str(&format!("  {}", comment));
+            let content = format!("  {}", comment);
+            output.push_str(&content);
+            line_length += content.len();
         }
     }
 
     for comment in &comments.inline {
-        output.push_str(&format!("  {}", comment));
+        let content = format!("  {}", comment);
+        output.push_str(&content);
+        line_length += content.len();
     }
 
     output.push('\n');
 
-    output
+    (output, line_length)
 }
 
 /// Format an import-from statement in multi-line format.
 fn format_multi_line(
     import_from: &ImportFromData,
     comments: &CommentSet,
     aliases: &[(AliasData, CommentSet)],
+    is_first: bool,
 ) -> String {
     let mut output = String::new();
 
+    if !is_first && !comments.atop.is_empty() {
+        output.push('\n');
+    }
     for comment in &comments.atop {
         output.push_str(comment);
         output.push('\n');

src/isort/mod.rs

@@ -168,10 +168,10 @@ fn normalize_imports(imports: Vec<AnnotatedImport>) -> ImportBlock {
                         })
                         .or_default();
                     for comment in atop {
-                        entry.atop.insert(comment.value);
+                        entry.atop.push(comment.value);
                     }
                     for comment in inline {
-                        entry.inline.insert(comment.value);
+                        entry.inline.push(comment.value);
                     }
                 }
 
@@ -202,10 +202,10 @@ fn normalize_imports(imports: Vec<AnnotatedImport>) -> ImportBlock {
                             .or_default()
                             .0;
                         for comment in atop {
-                            entry.atop.insert(comment.value);
+                            entry.atop.push(comment.value);
                         }
                         for comment in inline {
-                            entry.inline.insert(comment.value);
+                            entry.inline.push(comment.value);
                         }
                     } else {
                         let entry = &mut block
@@ -219,10 +219,10 @@ fn normalize_imports(imports: Vec<AnnotatedImport>) -> ImportBlock {
                             ))
                             .or_default();
                         for comment in atop {
-                            entry.atop.insert(comment.value);
+                            entry.atop.push(comment.value);
                         }
                         for comment in inline {
-                            entry.inline.insert(comment.value);
+                            entry.inline.push(comment.value);
                         }
                     }
                 }
@@ -241,10 +241,10 @@ fn normalize_imports(imports: Vec<AnnotatedImport>) -> ImportBlock {
                             })
                             .or_default();
                         for comment in alias.atop {
-                            entry.atop.insert(comment.value);
+                            entry.atop.push(comment.value);
                         }
                         for comment in alias.inline {
-                            entry.inline.insert(comment.value);
+                            entry.inline.push(comment.value);
                         }
                     } else {
                         let entry = block
@@ -257,11 +257,11 @@ fn normalize_imports(imports: Vec<AnnotatedImport>) -> ImportBlock {
                                 },
                             ))
                             .or_default();
-                        for comment in alias.atop {
-                            entry.atop.insert(comment.value);
-                        }
+                        entry
+                            .atop
+                            .extend(alias.atop.into_iter().map(|comment| comment.value));
                         for comment in alias.inline {
-                            entry.inline.insert(comment.value);
+                            entry.inline.push(comment.value);
                         }
                     }
                 }
@@ -414,7 +414,7 @@ pub fn format_imports(
     // Normalize imports (i.e., deduplicate, aggregate `from` imports).
     let block = normalize_imports(block);
 
-    println!("block = {:?}", block);
+    // println!("block = {:?}", block);
 
     // Categorize by type (e.g., first-party vs. third-party).
     let block_by_type = categorize_imports(
@@ -425,27 +425,39 @@ pub fn format_imports(
         extra_standard_library,
     );
 
-    // Generate replacement source code.
     let mut output = RopeBuilder::new();
-    let mut first_block = true;
+
+    // Generate replacement source code.
+    let mut is_first_block = true;
     for import_block in block_by_type.into_values() {
         let import_block = sort_imports(import_block);
 
         // Add a blank line between every section.
-        if !first_block {
+        if !is_first_block {
             output.append("\n");
         } else {
-            first_block = false;
+            is_first_block = false;
         }
 
+        let mut is_first_statement = true;
+
         // Format `StmtKind::Import` statements.
         for (alias, comments) in import_block.import.iter() {
-            format::format_import(&mut output, alias, comments);
+            format::format_import(&mut output, alias, comments, is_first_statement);
+            is_first_statement = false;
         }
 
         // Format `StmtKind::ImportFrom` statements.
         for (import_from, comments, aliases) in import_block.import_from.iter() {
-            format::format_import_from(&mut output, import_from, comments, aliases, line_length);
+            format::format_import_from(
+                &mut output,
+                import_from,
+                comments,
+                aliases,
+                line_length,
+                is_first_statement,
+            );
+            is_first_statement = false;
         }
     }
     output.finish().to_string()
@@ -463,14 +475,17 @@ mod tests {
     use crate::linter::test_path;
     use crate::Settings;
 
+    #[test_case(Path::new("add_newline_before_comments.py"))]
     #[test_case(Path::new("combine_import_froms.py"))]
     #[test_case(Path::new("comments.py"))]
     #[test_case(Path::new("deduplicate_imports.py"))]
     #[test_case(Path::new("fit_line_length.py"))]
+    #[test_case(Path::new("fit_line_length_comment.py"))]
     #[test_case(Path::new("import_from_after_import.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"))]
+    #[test_case(Path::new("preserve_comment_order.py"))]
     #[test_case(Path::new("preserve_indentation.py"))]
     #[test_case(Path::new("reorder_within_section.py"))]
     #[test_case(Path::new("separate_first_party_imports.py"))]

src/isort/snapshots/ruff__isort__tests__add_newline_before_comments.py.snap

@@ -0,0 +1,22 @@
+---
+source: src/isort/mod.rs
+expression: checks
+---
+- kind: UnsortedImports
+  location:
+    row: 1
+    column: 0
+  end_location:
+    row: 8
+    column: 0
+  fix:
+    patch:
+      content: "import os\n\n# This is a comment in the same section, so we need to add one newline.\nimport sys\n\nimport numpy as np\n\n# This is a comment, but it starts a new section, so we don't need to add a newline\n# before it.\nimport leading_prefix\n"
+      location:
+        row: 1
+        column: 0
+      end_location:
+        row: 8
+        column: 0
+    applied: false
+

src/isort/snapshots/ruff__isort__tests__comments.py.snap

@@ -11,7 +11,7 @@ expression: checks
     column: 0
   fix:
     patch:
-      content: "import B  # Comment 4\n# Comment 3a\n# Comment 3b\nimport C\nimport D\n# Comment 5\n# Comment 6\nfrom A import (\n    a,  # Comment 7  # Comment 9\n    b,  # Comment 10\n    c,  # Comment 11  # Comment 8\n)\n"
+      content: "import B  # Comment 4\n\n# Comment 3a\n# Comment 3b\nimport C\nimport D\n\n# Comment 5\n# Comment 6\nfrom A import (\n    a,  # Comment 7  # Comment 9\n    b,  # Comment 10\n    c,  # Comment 8  # Comment 11\n)\n"
       location:
         row: 3
         column: 0

src/isort/snapshots/ruff__isort__tests__fit_line_length_comment.py.snap

@@ -0,0 +1,22 @@
+---
+source: src/isort/mod.rs
+expression: checks
+---
+- kind: UnsortedImports
+  location:
+    row: 1
+    column: 0
+  end_location:
+    row: 5
+    column: 0
+  fix:
+    patch:
+      content: "import a\n\n# Don't take this comment into account when determining whether the next import can fit on one line.\nfrom b import c\nfrom d import (  # Do take this comment into account when determining whether the next import can fit on one line.\n    e,\n)\n"
+      location:
+        row: 1
+        column: 0
+      end_location:
+        row: 5
+        column: 0
+    applied: false
+

src/isort/snapshots/ruff__isort__tests__preserve_comment_order.py.snap

@@ -0,0 +1,22 @@
+---
+source: src/isort/mod.rs
+expression: checks
+---
+- kind: UnsortedImports
+  location:
+    row: 1
+    column: 0
+  end_location:
+    row: 12
+    column: 0
+  fix:
+    patch:
+      content: "import abc\nimport io\n\n# Old MacDonald had a farm,\n# EIEIO\n# And on his farm he had a cow,\n# EIEIO\n# With a moo-moo here and a moo-moo there\n# Here a moo, there a moo, everywhere moo-moo\n# Old MacDonald had a farm,\n# EIEIO\nfrom errno import EIO\n"
+      location:
+        row: 1
+        column: 0
+      end_location:
+        row: 12
+        column: 0
+    applied: false
+