Preserve comments when sorting imports (#749)

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/comments.py

@@ -0,0 +1,25 @@
+# Comment 1
+# Comment 2
+import D
+
+# Comment 3a
+import C
+
+# Comment 3b
+import C
+
+import B  # Comment 4
+
+# Comment 5
+
+# Comment 6
+from A import (
+    a,  # Comment 7
+    b,
+    c,  # Comment 8
+)
+from A import (
+    a,  # Comment 9
+    b,  # Comment 10
+    c,  # Comment 11
+)

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

resources/test/fixtures/isort/type_comments.py

@@ -0,0 +1,2 @@
+import A  # type: ignore
+from B import C  # type: ignore

src/isort/comments.rs

@@ -0,0 +1,40 @@
+use std::borrow::Cow;
+
+use rustpython_ast::Location;
+use rustpython_parser::lexer;
+use rustpython_parser::lexer::Tok;
+
+use crate::ast::helpers;
+use crate::ast::types::Range;
+use crate::SourceCodeLocator;
+
+#[derive(Debug)]
+pub struct Comment<'a> {
+    pub value: Cow<'a, str>,
+    pub location: Location,
+    pub end_location: Location,
+}
+
+/// Collect all comments in an import block.
+pub fn collect_comments<'a>(range: &Range, locator: &'a SourceCodeLocator) -> Vec<Comment<'a>> {
+    let contents = locator.slice_source_code_range(range);
+    lexer::make_tokenizer(&contents)
+        .flatten()
+        .filter_map(|(start, tok, end)| {
+            if matches!(tok, Tok::Comment) {
+                let start = helpers::to_absolute(&start, &range.location);
+                let end = helpers::to_absolute(&end, &range.location);
+                Some(Comment {
+                    value: locator.slice_source_code_range(&Range {
+                        location: start,
+                        end_location: end,
+                    }),
+                    location: start,
+                    end_location: end,
+                })
+            } else {
+                None
+            }
+        })
+        .collect()
+}

src/isort/format.rs

@@ -0,0 +1,182 @@
+use crate::isort::types::{AliasData, CommentSet, ImportFromData, Importable};
+
+// Hard-code four-space indentation for the imports themselves, to match Black.
+const INDENT: &str = "    ";
+
+// Guess a capacity to use for string allocation.
+const CAPACITY: usize = 200;
+
+/// Add a plain import statement to the `RopeBuilder`.
+pub fn format_import(alias: &AliasData, comments: &CommentSet, is_first: bool) -> String {
+    let mut output = String::with_capacity(CAPACITY);
+    if !is_first && !comments.atop.is_empty() {
+        output.push('\n');
+    }
+    for comment in &comments.atop {
+        output.push_str(comment);
+        output.push('\n');
+    }
+    if let Some(asname) = alias.asname {
+        output.push_str("import ");
+        output.push_str(alias.name);
+        output.push_str(" as ");
+        output.push_str(asname);
+    } else {
+        output.push_str("import ");
+        output.push_str(alias.name);
+    }
+    for comment in &comments.inline {
+        output.push_str("  ");
+        output.push_str(comment);
+    }
+    output.push('\n');
+    output
+}
+
+/// Add an import-from statement to the `RopeBuilder`.
+pub fn format_import_from(
+    import_from: &ImportFromData,
+    comments: &CommentSet,
+    aliases: &[(AliasData, CommentSet)],
+    line_length: &usize,
+    is_first: bool,
+) -> String {
+    // We can only inline if: (1) none of the aliases have atop comments, and (3)
+    // only the last alias (if any) has inline comments.
+    if aliases
+        .iter()
+        .all(|(_, CommentSet { atop, .. })| atop.is_empty())
+        && aliases
+            .iter()
+            .rev()
+            .skip(1)
+            .all(|(_, CommentSet { inline, .. })| inline.is_empty())
+    {
+        let (single_line, import_length) =
+            format_single_line(import_from, comments, aliases, is_first);
+        if import_length <= *line_length {
+            return single_line;
+        }
+    }
+
+    format_multi_line(import_from, comments, aliases, is_first)
+}
+
+/// Format an import-from statement in single-line format.
+///
+/// This method assumes that the output source code is syntactically valid.
+fn format_single_line(
+    import_from: &ImportFromData,
+    comments: &CommentSet,
+    aliases: &[(AliasData, CommentSet)],
+    is_first: bool,
+) -> (String, usize) {
+    let mut output = String::with_capacity(CAPACITY);
+    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');
+    }
+
+    let module_name = import_from.module_name();
+    output.push_str("from ");
+    output.push_str(&module_name);
+    output.push_str(" import ");
+    line_length += 5 + module_name.len() + 8;
+
+    for (index, (AliasData { name, asname }, comments)) in aliases.iter().enumerate() {
+        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(' ');
+            output.push(' ');
+            output.push_str(comment);
+            line_length += 2 + comment.len();
+        }
+    }
+
+    for comment in &comments.inline {
+        output.push(' ');
+        output.push(' ');
+        output.push_str(comment);
+        line_length += 2 + comment.len();
+    }
+
+    output.push('\n');
+
+    (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::with_capacity(CAPACITY);
+
+    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("from ");
+    output.push_str(&import_from.module_name());
+    output.push_str(" import ");
+    output.push('(');
+    for comment in &comments.inline {
+        output.push(' ');
+        output.push(' ');
+        output.push_str(comment);
+    }
+    output.push('\n');
+
+    for (AliasData { name, asname }, comments) in aliases {
+        for comment in &comments.atop {
+            output.push_str(INDENT);
+            output.push_str(comment);
+            output.push('\n');
+        }
+        output.push_str(INDENT);
+        if let Some(asname) = asname {
+            output.push_str(name);
+            output.push_str(" as ");
+            output.push_str(asname);
+        } else {
+            output.push_str(name);
+        }
+        output.push(',');
+
+        for comment in &comments.inline {
+            output.push(' ');
+            output.push(' ');
+            output.push_str(comment);
+        }
+        output.push('\n');
+    }
+
+    output.push(')');
+    output.push('\n');
+
+    output
+}

src/isort/mod.rs

@@ -1,64 +1,266 @@
 use std::collections::{BTreeMap, BTreeSet};
 use std::path::PathBuf;
 
-use fnv::FnvHashSet;
+use fnv::FnvHashMap;
 use itertools::Itertools;
 use ropey::RopeBuilder;
 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::types::{AliasData, ImportBlock, ImportFromData, Importable, OrderedImportBlock};
+use crate::isort::types::{
+    AliasData, CommentSet, ImportBlock, ImportFromData, Importable, OrderedImportBlock,
+};
 
 mod categorize;
+mod comments;
+pub mod format;
 pub mod plugins;
 pub mod settings;
 mod sorting;
 pub mod track;
 mod types;
 
-// Hard-code four-space indentation for the imports themselves, to match Black.
-const INDENT: &str = "    ";
+#[derive(Debug)]
+pub struct AnnotatedAliasData<'a> {
+    pub name: &'a str,
+    pub asname: &'a Option<String>,
+    pub atop: Vec<Comment<'a>>,
+    pub inline: Vec<Comment<'a>>,
+}
+#[derive(Debug)]
+pub enum AnnotatedImport<'a> {
+    Import {
+        names: Vec<AliasData<'a>>,
+        atop: Vec<Comment<'a>>,
+        inline: Vec<Comment<'a>>,
+    },
+    ImportFrom {
+        module: &'a Option<String>,
+        names: Vec<AnnotatedAliasData<'a>>,
+        level: &'a Option<usize>,
+        atop: Vec<Comment<'a>>,
+        inline: Vec<Comment<'a>>,
+    },
+}
 
-fn normalize_imports<'a>(imports: &'a [&'a Stmt]) -> ImportBlock<'a> {
-    let mut block: ImportBlock = Default::default();
+fn annotate_imports<'a>(
+    imports: &'a [&'a Stmt],
+    comments: Vec<Comment<'a>>,
+) -> Vec<AnnotatedImport<'a>> {
+    let mut annotated = vec![];
+    let mut comments_iter = comments.into_iter().peekable();
     for import in imports {
         match &import.node {
             StmtKind::Import { names } => {
-                for name in names {
-                    block.import.insert(AliasData {
-                        name: &name.node.name,
-                        asname: &name.node.asname,
-                    });
+                // Find comments above.
+                let mut atop = vec![];
+                while let Some(comment) =
+                    comments_iter.next_if(|comment| comment.location.row() < import.location.row())
+                {
+                    atop.push(comment);
                 }
+
+                // Find comments inline.
+                let mut inline = vec![];
+                while let Some(comment) = comments_iter.next_if(|comment| {
+                    comment.end_location.row() == import.end_location.unwrap().row()
+                }) {
+                    inline.push(comment);
+                }
+
+                annotated.push(AnnotatedImport::Import {
+                    names: names
+                        .iter()
+                        .map(|alias| AliasData {
+                            name: &alias.node.name,
+                            asname: &alias.node.asname,
+                        })
+                        .collect(),
+                    atop,
+                    inline,
+                });
             }
             StmtKind::ImportFrom {
                 module,
                 names,
                 level,
             } => {
-                for name in names {
-                    if name.node.asname.is_none() {
-                        block
+                // Find comments above.
+                let mut atop = vec![];
+                while let Some(comment) =
+                    comments_iter.next_if(|comment| comment.location.row() < import.location.row())
+                {
+                    atop.push(comment);
+                }
+
+                // Find comments inline.
+                let mut inline = vec![];
+                while let Some(comment) =
+                    comments_iter.next_if(|comment| comment.location.row() == import.location.row())
+                {
+                    inline.push(comment);
+                }
+
+                // Capture names.
+                let mut aliases = vec![];
+                for alias in names {
+                    // Find comments above.
+                    let mut alias_atop = vec![];
+                    while let Some(comment) = comments_iter
+                        .next_if(|comment| comment.location.row() < alias.location.row())
+                    {
+                        alias_atop.push(comment);
+                    }
+
+                    // Find comments inline.
+                    let mut alias_inline = vec![];
+                    while let Some(comment) = comments_iter.next_if(|comment| {
+                        comment.end_location.row() == alias.end_location.unwrap().row()
+                    }) {
+                        alias_inline.push(comment);
+                    }
+
+                    aliases.push(AnnotatedAliasData {
+                        name: &alias.node.name,
+                        asname: &alias.node.asname,
+                        atop: alias_atop,
+                        inline: alias_inline,
+                    })
+                }
+
+                annotated.push(AnnotatedImport::ImportFrom {
+                    module,
+                    names: aliases,
+                    level,
+                    atop,
+                    inline,
+                });
+            }
+            _ => unreachable!("Expected StmtKind::Import | StmtKind::ImportFrom"),
+        }
+    }
+    annotated
+}
+
+fn normalize_imports(imports: Vec<AnnotatedImport>) -> ImportBlock {
+    let mut block: ImportBlock = Default::default();
+    for import in imports {
+        match import {
+            AnnotatedImport::Import {
+                names,
+                atop,
+                inline,
+            } => {
+                // Associate the comments with the first alias (best effort).
+                if let Some(name) = names.first() {
+                    let entry = block
+                        .import
+                        .entry(AliasData {
+                            name: name.name,
+                            asname: name.asname,
+                        })
+                        .or_default();
+                    for comment in atop {
+                        entry.atop.push(comment.value);
+                    }
+                    for comment in inline {
+                        entry.inline.push(comment.value);
+                    }
+                }
+
+                // Create an entry for every alias.
+                for name in &names {
+                    block
+                        .import
+                        .entry(AliasData {
+                            name: name.name,
+                            asname: name.asname,
+                        })
+                        .or_default();
+                }
+            }
+            AnnotatedImport::ImportFrom {
+                module,
+                names,
+                level,
+                atop,
+                inline,
+            } => {
+                // Associate the comments with the first alias (best effort).
+                if let Some(alias) = names.first() {
+                    if alias.asname.is_none() {
+                        let entry = &mut block
                             .import_from
                             .entry(ImportFromData { module, level })
                             .or_default()
-                            .insert(AliasData {
-                                name: &name.node.name,
-                                asname: &name.node.asname,
-                            });
+                            .0;
+                        for comment in atop {
+                            entry.atop.push(comment.value);
+                        }
+                        for comment in inline {
+                            entry.inline.push(comment.value);
+                        }
                     } else {
-                        block.import_from_as.insert((
-                            ImportFromData { module, level },
-                            AliasData {
-                                name: &name.node.name,
-                                asname: &name.node.asname,
-                            },
-                        ));
+                        let entry = &mut block
+                            .import_from_as
+                            .entry((
+                                ImportFromData { module, level },
+                                AliasData {
+                                    name: alias.name,
+                                    asname: alias.asname,
+                                },
+                            ))
+                            .or_default();
+                        for comment in atop {
+                            entry.atop.push(comment.value);
+                        }
+                        for comment in inline {
+                            entry.inline.push(comment.value);
+                        }
+                    }
+                }
+
+                // Create an entry for every alias.
+                for alias in names {
+                    if alias.asname.is_none() {
+                        let entry = block
+                            .import_from
+                            .entry(ImportFromData { module, level })
+                            .or_default()
+                            .1
+                            .entry(AliasData {
+                                name: alias.name,
+                                asname: alias.asname,
+                            })
+                            .or_default();
+                        for comment in alias.atop {
+                            entry.atop.push(comment.value);
+                        }
+                        for comment in alias.inline {
+                            entry.inline.push(comment.value);
+                        }
+                    } else {
+                        let entry = block
+                            .import_from_as
+                            .entry((
+                                ImportFromData { module, level },
+                                AliasData {
+                                    name: alias.name,
+                                    asname: alias.asname,
+                                },
+                            ))
+                            .or_default();
+                        entry
+                            .atop
+                            .extend(alias.atop.into_iter().map(|comment| comment.value));
+                        for comment in alias.inline {
+                            entry.inline.push(comment.value);
+                        }
                     }
                 }
             }
-            _ => unreachable!("Expected StmtKind::Import | StmtKind::ImportFrom"),
         }
     }
     block
@@ -73,7 +275,7 @@ fn categorize_imports<'a>(
 ) -> BTreeMap<ImportType, ImportBlock<'a>> {
     let mut block_by_type: BTreeMap<ImportType, ImportBlock> = Default::default();
     // Categorize `StmtKind::Import`.
-    for alias in block.import {
+    for (alias, comments) in block.import {
         let import_type = categorize(
             &alias.module_base(),
             &None,
@@ -86,7 +288,7 @@ fn categorize_imports<'a>(
             .entry(import_type)
             .or_default()
             .import
-            .insert(alias);
+            .insert(alias, comments);
     }
     // Categorize `StmtKind::ImportFrom` (without re-export).
     for (import_from, aliases) in block.import_from {
@@ -105,7 +307,7 @@ fn categorize_imports<'a>(
             .insert(import_from, aliases);
     }
     // Categorize `StmtKind::ImportFrom` (with re-export).
-    for (import_from, alias) in block.import_from_as {
+    for ((import_from, alias), comments) in block.import_from_as {
         let classification = categorize(
             &import_from.module_base(),
             import_from.level,
@@ -118,7 +320,7 @@ fn categorize_imports<'a>(
             .entry(classification)
             .or_default()
             .import_from_as
-            .insert((import_from, alias));
+            .insert((import_from, alias), comments);
     }
     block_by_type
 }
@@ -131,7 +333,7 @@ fn sort_imports(block: ImportBlock) -> OrderedImportBlock {
         block
             .import
             .into_iter()
-            .sorted_by_cached_key(|alias| module_key(alias.name, alias.asname)),
+            .sorted_by_cached_key(|(alias, _)| module_key(alias.name, alias.asname)),
     );
 
     // Sort `StmtKind::ImportFrom`.
@@ -145,19 +347,37 @@ fn sort_imports(block: ImportBlock) -> OrderedImportBlock {
                 block
                     .import_from_as
                     .into_iter()
-                    .map(|(import_from, alias)| (import_from, FnvHashSet::from_iter([alias]))),
+                    .map(|((import_from, alias), comments)| {
+                        (
+                            import_from,
+                            (
+                                CommentSet {
+                                    atop: comments.atop,
+                                    inline: Default::default(),
+                                },
+                                FnvHashMap::from_iter([(
+                                    alias,
+                                    CommentSet {
+                                        atop: Default::default(),
+                                        inline: comments.inline,
+                                    },
+                                )]),
+                            ),
+                        )
+                    }),
             )
-            .map(|(import_from, aliases)| {
+            .map(|(import_from, (comments, aliases))| {
                 // Within each `StmtKind::ImportFrom`, sort the members.
                 (
                     import_from,
+                    comments,
                     aliases
                         .into_iter()
-                        .sorted_by_cached_key(|alias| member_key(alias.name, alias.asname))
-                        .collect::<Vec<AliasData>>(),
+                        .sorted_by_cached_key(|(alias, _)| member_key(alias.name, alias.asname))
+                        .collect::<Vec<(AliasData, CommentSet)>>(),
                 )
             })
-            .sorted_by_cached_key(|(import_from, aliases)| {
+            .sorted_by_cached_key(|(import_from, _, aliases)| {
                 // Sort each `StmtKind::ImportFrom` by module key, breaking ties based on
                 // members.
                 (
@@ -167,7 +387,7 @@ fn sort_imports(block: ImportBlock) -> OrderedImportBlock {
                         .map(|module| module_key(module, &None)),
                     aliases
                         .first()
-                        .map(|alias| member_key(alias.name, alias.asname)),
+                        .map(|(alias, _)| member_key(alias.name, alias.asname)),
                 )
             }),
     );
@@ -176,15 +396,18 @@ fn sort_imports(block: ImportBlock) -> OrderedImportBlock {
 }
 
 pub fn format_imports(
-    block: Vec<&Stmt>,
+    block: &[&Stmt],
+    comments: Vec<Comment>,
     line_length: &usize,
     src: &[PathBuf],
     known_first_party: &BTreeSet<String>,
     known_third_party: &BTreeSet<String>,
     extra_standard_library: &BTreeSet<String>,
 ) -> String {
+    let block = annotate_imports(block, comments);
+
     // Normalize imports (i.e., deduplicate, aggregate `from` imports).
-    let block = normalize_imports(&block);
+    let block = normalize_imports(block);
 
     // Categorize by type (e.g., first-party vs. third-party).
     let block_by_type = categorize_imports(
@@ -195,81 +418,38 @@ 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 AliasData { name, asname } in import_block.import.iter() {
-            if let Some(asname) = asname {
-                output.append(&format!("import {} as {}\n", name, asname));
-            } else {
-                output.append(&format!("import {}\n", name));
-            }
+        for (alias, comments) in import_block.import.iter() {
+            output.append(&format::format_import(alias, comments, is_first_statement));
+            is_first_statement = false;
         }
 
         // Format `StmtKind::ImportFrom` statements.
-        for (import_from, aliases) in import_block.import_from.iter() {
-            let prelude: String = format!("from {} import ", import_from.module_name());
-            let members: Vec<String> = aliases
-                .iter()
-                .map(|AliasData { name, asname }| {
-                    if let Some(asname) = asname {
-                        format!("{} as {}", name, asname)
-                    } else {
-                        name.to_string()
-                    }
-                })
-                .collect();
-
-            // Can we fit the import on a single line?
-            let expected_len: usize =
-                    // `from base import `
-                    prelude.len()
-                        // `member( as alias)?`
-                        + members.iter().map(|part| part.len()).sum::<usize>()
-                        // `, `
-                        + 2 * (members.len() - 1);
-
-            if expected_len <= *line_length {
-                // `from base import `
-                output.append(&prelude);
-                // `member( as alias)?(, )?`
-                for (index, part) in members.into_iter().enumerate() {
-                    if index > 0 {
-                        output.append(", ");
-                    }
-                    output.append(&part);
-                }
-                // `\n`
-                output.append("\n");
-            } else {
-                // `from base import (\n`
-                output.append(&prelude);
-                output.append("(");
-                output.append("\n");
-
-                // `    member( as alias)?,\n`
-                for part in members {
-                    output.append(INDENT);
-                    output.append(&part);
-                    output.append(",");
-                    output.append("\n");
-                }
-
-                // `)\n`
-                output.append(")");
-                output.append("\n");
-            }
+        for (import_from, comments, aliases) in import_block.import_from.iter() {
+            output.append(&format::format_import_from(
+                import_from,
+                comments,
+                aliases,
+                line_length,
+                is_first_statement,
+            ));
+            is_first_statement = false;
         }
     }
     output.finish().to_string()
@@ -287,13 +467,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"))]
@@ -303,6 +487,7 @@ mod tests {
     #[test_case(Path::new("skip.py"))]
     #[test_case(Path::new("sort_similar_imports.py"))]
     #[test_case(Path::new("trailing_suffix.py"))]
+    #[test_case(Path::new("type_comments.py"))]
     fn isort(path: &Path) -> Result<()> {
         let snapshot = format!("{}", path.to_string_lossy());
         let mut checks = test_path(

src/isort/plugins.rs

@@ -5,7 +5,7 @@ use crate::ast::types::Range;
 use crate::autofix::{fixer, Fix};
 use crate::checks::CheckKind;
 use crate::docstrings::helpers::leading_space;
-use crate::isort::format_imports;
+use crate::isort::{comments, format_imports};
 use crate::{Check, Settings, SourceCodeLocator};
 
 fn extract_range(body: &[&Stmt]) -> Range {
@@ -44,7 +44,15 @@ fn match_trailing_content(body: &[&Stmt], locator: &SourceCodeLocator) -> bool {
         end_location: Location::new(end_location.row() + 1, 0),
     };
     let suffix = locator.slice_source_code_range(&range);
-    suffix.chars().any(|char| !char.is_whitespace())
+    for char in suffix.chars() {
+        if char == '#' {
+            return false;
+        }
+        if !char.is_whitespace() {
+            return true;
+        }
+    }
+    false
 }
 
 /// I001
@@ -57,13 +65,23 @@ pub fn check_imports(
     let range = extract_range(&body);
     let indentation = extract_indentation(&body, locator);
 
+    // Extract comments. Take care to grab any inline comments from the last line.
+    let comments = comments::collect_comments(
+        &Range {
+            location: range.location,
+            end_location: Location::new(range.end_location.row() + 1, 0),
+        },
+        locator,
+    );
+
     // Special-cases: there's leading or trailing content in the import block.
     let has_leading_content = match_leading_content(&body, locator);
     let has_trailing_content = match_trailing_content(&body, locator);
 
     // Generate the sorted import block.
     let expected = format_imports(
-        body,
+        &body,
+        comments,
         &(settings.line_length - indentation.len()),
         &settings.src,
         &settings.isort.known_first_party,

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

@@ -0,0 +1,22 @@
+---
+source: src/isort/mod.rs
+expression: checks
+---
+- kind: UnsortedImports
+  location:
+    row: 3
+    column: 0
+  end_location:
+    row: 26
+    column: 0
+  fix:
+    patch:
+      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
+      end_location:
+        row: 26
+        column: 0
+    applied: false
+

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
+

src/isort/snapshots/ruff__isort__tests__type_comments.py.snap

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

src/isort/types.rs

@@ -1,4 +1,6 @@
-use fnv::{FnvHashMap, FnvHashSet};
+use std::borrow::Cow;
+
+use fnv::FnvHashMap;
 
 #[derive(Debug, Hash, Ord, PartialOrd, Eq, PartialEq)]
 pub struct ImportFromData<'a> {
@@ -12,6 +14,12 @@ pub struct AliasData<'a> {
     pub asname: &'a Option<String>,
 }
 
+#[derive(Debug, Default)]
+pub struct CommentSet<'a> {
+    pub atop: Vec<Cow<'a, str>>,
+    pub inline: Vec<Cow<'a, str>>,
+}
+
 pub trait Importable {
     fn module_name(&self) -> String;
     fn module_base(&self) -> String;
@@ -50,17 +58,24 @@ impl Importable for ImportFromData<'_> {
 pub struct ImportBlock<'a> {
     // Set of (name, asname), used to track regular imports.
     // Ex) `import module`
-    pub import: FnvHashSet<AliasData<'a>>,
+    pub import: FnvHashMap<AliasData<'a>, CommentSet<'a>>,
     // Map from (module, level) to `AliasData`, used to track 'from' imports.
     // Ex) `from module import member`
-    pub import_from: FnvHashMap<ImportFromData<'a>, FnvHashSet<AliasData<'a>>>,
+    pub import_from:
+        FnvHashMap<ImportFromData<'a>, (CommentSet<'a>, FnvHashMap<AliasData<'a>, CommentSet<'a>>)>,
     // Set of (module, level, name, asname), used to track re-exported 'from' imports.
     // Ex) `from module import member as member`
-    pub import_from_as: FnvHashSet<(ImportFromData<'a>, AliasData<'a>)>,
+    pub import_from_as: FnvHashMap<(ImportFromData<'a>, AliasData<'a>), CommentSet<'a>>,
 }
 
+type AliasDataWithComments<'a> = (AliasData<'a>, CommentSet<'a>);
+
 #[derive(Debug, Default)]
 pub struct OrderedImportBlock<'a> {
-    pub import: Vec<AliasData<'a>>,
-    pub import_from: Vec<(ImportFromData<'a>, Vec<AliasData<'a>>)>,
+    pub import: Vec<AliasDataWithComments<'a>>,
+    pub import_from: Vec<(
+        ImportFromData<'a>,
+        CommentSet<'a>,
+        Vec<AliasDataWithComments<'a>>,
+    )>,
 }