Preserve trailing inline comments on import-from statements

crates/ruff_linter/resources/test/fixtures/isort/trailing_comment.py

@@ -0,0 +1,55 @@
+from mylib import (
+    MyClient,
+    MyMgmtClient,
+)  # some comment
+
+pass
+
+from mylib import (
+    MyClient,
+    MyMgmtClient,
+); from foo import (
+    bar
+)# some comment
+
+pass
+
+from foo import (
+    bar
+)
+from mylib import (
+    MyClient,
+    MyMgmtClient,
+    # some comment
+)
+
+pass
+
+from mylib import (
+    MyClient,
+    # some comment
+)
+
+pass
+
+from mylib import (
+    MyClient
+    # some comment
+)
+
+pass
+
+from mylib import (
+    # some comment
+    MyClient
+)
+
+pass
+
+# a
+from mylib import ( # b
+    # c
+    MyClient # d
+    # e
+) # f
+

crates/ruff_linter/src/rules/isort/annotate.rs

@@ -87,7 +87,7 @@ pub(crate) fn annotate_imports<'a>(
                     }
 
                     // Capture names.
-                    let aliases = names
+                    let mut aliases: Vec<_> = names
                         .iter()
                         .map(|alias| {
                             // Find comments above.
@@ -112,10 +112,40 @@ pub(crate) fn annotate_imports<'a>(
                                 asname: alias.asname.as_ref().map(|asname| locator.slice(asname)),
                                 atop: alias_atop,
                                 inline: alias_inline,
+                                trailing: vec![],
                             }
                         })
                         .collect();
 
+                    // Capture trailing comments on the _last_ alias, as in:
+                    // ```python
+                    // from foo import (
+                    //     bar,
+                    //     # noqa
+                    // )
+                    // ```
+                    if let Some(last_alias) = aliases.last_mut() {
+                        while let Some(comment) =
+                            comments_iter.next_if(|comment| comment.start() < import.end())
+                        {
+                            last_alias.trailing.push(comment);
+                        }
+                    }
+
+                    // Capture trailing comments, as in:
+                    // ```python
+                    // from foo import (
+                    //     bar,
+                    // )  # noqa
+                    // ```
+                    let mut trailing = vec![];
+                    let import_line_end = locator.line_end(import.end());
+                    while let Some(comment) =
+                        comments_iter.next_if(|comment| comment.start() < import_line_end)
+                    {
+                        trailing.push(comment);
+                    }
+
                     AnnotatedImport::ImportFrom {
                         module: module.as_ref().map(|module| locator.slice(module)),
                         names: aliases,
@@ -127,6 +157,7 @@ pub(crate) fn annotate_imports<'a>(
                         },
                         atop,
                         inline,
+                        trailing,
                     }
                 }
                 _ => panic!("Expected Stmt::Import | Stmt::ImportFrom"),

crates/ruff_linter/src/rules/isort/format.rs

@@ -1,16 +1,18 @@
+use std::borrow::Cow;
+
 use ruff_python_codegen::Stylist;
 
 use crate::line_width::{LineLength, LineWidthBuilder};
 
-use super::types::{AliasData, CommentSet, ImportFromData, Importable};
+use super::types::{AliasData, ImportCommentSet, ImportFromCommentSet, ImportFromData, Importable};
 
 // Guess a capacity to use for string allocation.
 const CAPACITY: usize = 200;
 
 /// Add a plain import statement to the [`RopeBuilder`].
 pub(crate) fn format_import(
     alias: &AliasData,
-    comments: &CommentSet,
+    comments: &ImportCommentSet,
     is_first: bool,
     stylist: &Stylist,
 ) -> String {
@@ -43,8 +45,8 @@ pub(crate) fn format_import(
 #[allow(clippy::too_many_arguments)]
 pub(crate) fn format_import_from(
     import_from: &ImportFromData,
-    comments: &CommentSet,
-    aliases: &[(AliasData, CommentSet)],
+    comments: &ImportFromCommentSet,
+    aliases: &[(AliasData, ImportFromCommentSet)],
     line_length: LineLength,
     indentation_width: LineWidthBuilder,
     stylist: &Stylist,
@@ -68,12 +70,19 @@ pub(crate) fn format_import_from(
         return single_line;
     }
 
-    // We can only inline if none of the aliases have atop or inline comments.
+    // We can only inline if none of the aliases have comments.
     if !trailing_comma
         && (aliases.len() == 1
-            || aliases
-                .iter()
-                .all(|(_, CommentSet { atop, inline })| atop.is_empty() && inline.is_empty()))
+            || aliases.iter().all(
+                |(
+                    _,
+                    ImportFromCommentSet {
+                        atop,
+                        inline,
+                        trailing,
+                    },
+                )| atop.is_empty() && inline.is_empty() && trailing.is_empty(),
+            ))
         && (!force_wrap_aliases
             || aliases.len() == 1
             || aliases.iter().all(|(alias, _)| alias.asname.is_none()))
@@ -99,8 +108,8 @@ pub(crate) fn format_import_from(
 /// This method assumes that the output source code is syntactically valid.
 fn format_single_line(
     import_from: &ImportFromData,
-    comments: &CommentSet,
-    aliases: &[(AliasData, CommentSet)],
+    comments: &ImportFromCommentSet,
+    aliases: &[(AliasData, ImportFromCommentSet)],
     is_first: bool,
     stylist: &Stylist,
     indentation_width: LineWidthBuilder,
@@ -136,16 +145,39 @@ fn format_single_line(
             output.push_str(", ");
             line_width = line_width.add_width(2);
         }
+    }
+
+    for comment in &comments.inline {
+        output.push(' ');
+        output.push(' ');
+        output.push_str(comment);
+        line_width = line_width.add_width(2).add_str(comment);
+    }
+
+    for (_, comments) in aliases.iter() {
+        for comment in &comments.atop {
+            output.push(' ');
+            output.push(' ');
+            output.push_str(comment);
+            line_width = line_width.add_width(2).add_str(comment);
+        }
 
         for comment in &comments.inline {
             output.push(' ');
             output.push(' ');
             output.push_str(comment);
             line_width = line_width.add_width(2).add_str(comment);
         }
+
+        for comment in &comments.trailing {
+            output.push(' ');
+            output.push(' ');
+            output.push_str(comment);
+            line_width = line_width.add_width(2).add_str(comment);
+        }
     }
 
-    for comment in &comments.inline {
+    for comment in &comments.trailing {
         output.push(' ');
         output.push(' ');
         output.push_str(comment);
@@ -160,8 +192,8 @@ fn format_single_line(
 /// Format an import-from statement in multi-line format.
 fn format_multi_line(
     import_from: &ImportFromData,
-    comments: &CommentSet,
-    aliases: &[(AliasData, CommentSet)],
+    comments: &ImportFromCommentSet,
+    aliases: &[(AliasData, ImportFromCommentSet)],
     is_first: bool,
     stylist: &Stylist,
 ) -> String {
@@ -208,9 +240,20 @@ fn format_multi_line(
             output.push_str(comment);
         }
         output.push_str(&stylist.line_ending());
+
+        for comment in &comments.trailing {
+            output.push_str(stylist.indentation());
+            output.push_str(comment);
+            output.push_str(&stylist.line_ending());
+        }
     }
 
     output.push(')');
+
+    for comment in &comments.trailing {
+        output.push_str("  ");
+        output.push_str(comment);
+    }
     output.push_str(&stylist.line_ending());
 
     output

crates/ruff_linter/src/rules/isort/mod.rs

@@ -41,6 +41,7 @@ pub(crate) struct AnnotatedAliasData<'a> {
     pub(crate) asname: Option<&'a str>,
     pub(crate) atop: Vec<Comment<'a>>,
     pub(crate) inline: Vec<Comment<'a>>,
+    pub(crate) trailing: Vec<Comment<'a>>,
 }
 
 #[derive(Debug)]
@@ -56,6 +57,7 @@ pub(crate) enum AnnotatedImport<'a> {
         level: u32,
         atop: Vec<Comment<'a>>,
         inline: Vec<Comment<'a>>,
+        trailing: Vec<Comment<'a>>,
         trailing_comma: TrailingComma,
     },
 }
@@ -84,6 +86,8 @@ pub(crate) fn format_imports(
         tokens,
     );
 
+    println!("{:?}", block);
+
     // Normalize imports (i.e., deduplicate, aggregate `from` imports).
     let block = normalize_imports(block, settings);
 
@@ -342,6 +346,7 @@ mod tests {
     #[test_case(Path::new("sort_similar_imports.py"))]
     #[test_case(Path::new("split.py"))]
     #[test_case(Path::new("star_before_others.py"))]
+    #[test_case(Path::new("trailing_comment.py"))]
     #[test_case(Path::new("trailing_suffix.py"))]
     #[test_case(Path::new("two_space.py"))]
     #[test_case(Path::new("type_comments.py"))]

crates/ruff_linter/src/rules/isort/normalize.rs

@@ -48,6 +48,7 @@ pub(crate) fn normalize_imports<'a>(
                 level,
                 atop,
                 inline,
+                trailing,
                 trailing_comma,
             } => {
                 // Whether to track each member of the import as a separate entry.
@@ -80,10 +81,13 @@ pub(crate) fn normalize_imports<'a>(
                             }
                         }
 
-                        // Replicate the inline comments onto every member.
+                        // Replicate the inline (and after) comments onto every member.
                         for comment in &inline {
                             import_from.comments.inline.push(comment.value.clone());
                         }
+                        for comment in &trailing {
+                            import_from.comments.trailing.push(comment.value.clone());
+                        }
                     }
                 } else {
                     if let Some(alias) = names.first() {
@@ -113,10 +117,12 @@ pub(crate) fn normalize_imports<'a>(
                         for comment in atop {
                             import_from.comments.atop.push(comment.value);
                         }
-
                         for comment in inline {
                             import_from.comments.inline.push(comment.value);
                         }
+                        for comment in trailing {
+                            import_from.comments.trailing.push(comment.value);
+                        }
                     }
                 }
 
@@ -161,6 +167,9 @@ pub(crate) fn normalize_imports<'a>(
                     for comment in alias.inline {
                         comment_set.inline.push(comment.value);
                     }
+                    for comment in alias.trailing {
+                        comment_set.trailing.push(comment.value);
+                    }
 
                     // Propagate trailing commas.
                     if !isolate_aliases && matches!(trailing_comma, TrailingComma::Present) {

crates/ruff_linter/src/rules/isort/order.rs

@@ -5,7 +5,7 @@ use itertools::Itertools;
 use super::settings::Settings;
 use super::sorting::{MemberKey, ModuleKey};
 use super::types::EitherImport::{self, Import, ImportFrom};
-use super::types::{AliasData, CommentSet, ImportBlock, ImportFromStatement};
+use super::types::{AliasData, ImportBlock, ImportFromCommentSet, ImportFromStatement};
 
 pub(crate) fn order_imports<'a>(
     block: ImportBlock<'a>,
@@ -49,7 +49,7 @@ pub(crate) fn order_imports<'a>(
                             .sorted_by_cached_key(|(alias, _)| {
                                 MemberKey::from_member(alias.name, alias.asname, settings)
                             })
-                            .collect::<Vec<(AliasData, CommentSet)>>(),
+                            .collect::<Vec<(AliasData, ImportFromCommentSet)>>(),
                     )
                 },
             );

crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__as_imports_comments.py.snap

@@ -36,10 +36,8 @@ as_imports_comments.py:1:1: I001 [*] Import block is un-sorted or un-formatted
 12    |-
 13    |-from bop import (  # Comment on `bop`
 14    |-    Member  # Comment on `Member`
-   4  |+from baz import Member as Alias  # Comment on `Alias`  # Comment on `baz`
-   5  |+from bop import Member  # Comment on `Member`  # Comment on `bop`
+   4  |+from baz import Member as Alias  # Comment on `baz`  # Comment on `Alias`
+   5  |+from bop import Member  # Comment on `bop`  # Comment on `Member`
    6  |+from foo import (  # Comment on `foo`
    7  |+    Member as Alias,  # Comment on `Alias`
 15 8  | )
-
-