Preserve trailing inline comments on import-from statements

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

@@ -0,0 +1,13 @@
+from mylib import (
+    MyClient,
+    MyMgmtClient,
+)  # some comment
+
+pass
+
+from mylib import (
+    MyClient,
+    MyMgmtClient,
+); from foo import (
+    bar
+)# some comment

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

@@ -16,6 +16,7 @@ pub(crate) fn annotate_imports<'a>(
     split_on_trailing_comma: bool,
     tokens: &Tokens,
 ) -> Vec<AnnotatedImport<'a>> {
+    println!("comments: {:?}", comments);
     let mut comments_iter = comments.into_iter().peekable();
 
     imports
@@ -116,6 +117,20 @@ pub(crate) fn annotate_imports<'a>(
                         })
                         .collect();
 
+                    // Capture trailing comments, as in:
+                    // ```python
+                    // from foo import (
+                    //     bar,
+                    // )  # noqa
+                    // ```
+                    let mut after = vec![];
+                    let import_line_end = locator.line_end(import.end());
+                    while let Some(comment) = comments_iter.next_if(|comment| {
+                        import.end() <= comment.start() && comment.start() < import_line_end
+                    }) {
+                        after.push(comment);
+                    }
+
                     AnnotatedImport::ImportFrom {
                         module: module.as_ref().map(|module| locator.slice(module)),
                         names: aliases,
@@ -127,6 +142,7 @@ pub(crate) fn annotate_imports<'a>(
                         },
                         atop,
                         inline,
+                        after,
                     }
                 }
                 _ => panic!("Expected Stmt::Import | Stmt::ImportFrom"),

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

@@ -1,4 +1,5 @@
 use ruff_python_codegen::Stylist;
+use std::borrow::Cow;
 
 use crate::line_width::{LineLength, LineWidthBuilder};
 
@@ -44,6 +45,7 @@ pub(crate) fn format_import(
 pub(crate) fn format_import_from(
     import_from: &ImportFromData,
     comments: &CommentSet,
+    trailing_comments: &[Cow<'_, str>],
     aliases: &[(AliasData, CommentSet)],
     line_length: LineLength,
     indentation_width: LineWidthBuilder,
@@ -70,6 +72,7 @@ pub(crate) fn format_import_from(
 
     // We can only inline if none of the aliases have atop or inline comments.
     if !trailing_comma
+        && trailing_comments.is_empty()
         && (aliases.len() == 1
             || aliases
                 .iter()
@@ -91,7 +94,14 @@ pub(crate) fn format_import_from(
         }
     }
 
-    format_multi_line(import_from, comments, aliases, is_first, stylist)
+    format_multi_line(
+        import_from,
+        comments,
+        trailing_comments,
+        aliases,
+        is_first,
+        stylist,
+    )
 }
 
 /// Format an import-from statement in single-line format.
@@ -161,6 +171,7 @@ fn format_single_line(
 fn format_multi_line(
     import_from: &ImportFromData,
     comments: &CommentSet,
+    trailing_comments: &[Cow<'_, str>],
     aliases: &[(AliasData, CommentSet)],
     is_first: bool,
     stylist: &Stylist,
@@ -211,6 +222,11 @@ fn format_multi_line(
     }
 
     output.push(')');
+
+    for comment in trailing_comments {
+        output.push_str("  ");
+        output.push_str(comment);
+    }
     output.push_str(&stylist.line_ending());
 
     output

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

@@ -56,6 +56,7 @@ pub(crate) enum AnnotatedImport<'a> {
         level: u32,
         atop: Vec<Comment<'a>>,
         inline: Vec<Comment<'a>>,
+        after: Vec<Comment<'a>>,
         trailing_comma: TrailingComma,
     },
 }
@@ -83,6 +84,7 @@ pub(crate) fn format_imports(
         settings.split_on_trailing_comma,
         tokens,
     );
+    println!("block: {:?}", block);
 
     // Normalize imports (i.e., deduplicate, aggregate `from` imports).
     let block = normalize_imports(block, settings);
@@ -236,7 +238,7 @@ fn format_import_block(
                     }
                 }
 
-                ImportFrom((import_from, comments, trailing_comma, aliases)) => {
+                ImportFrom((import_from, comments, trailing_comments, trailing_comma, aliases)) => {
                     // Add a blank lines between direct and from imports.
                     if !settings.from_first
                         && lines_between_types > 0
@@ -253,6 +255,7 @@ fn format_import_block(
                     output.push_str(&format::format_import_from(
                         &import_from,
                         &comments,
+                        &trailing_comments,
                         &aliases,
                         line_length,
                         indentation_width,
@@ -342,6 +345,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,
+                after,
                 trailing_comma,
             } => {
                 // Whether to track each member of the import as a separate entry.
@@ -80,10 +81,14 @@ 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 &after {
+                            import_from.comments.inline.push(comment.value.clone());
+                        }
                     }
                 } else {
                     if let Some(alias) = names.first() {
@@ -117,6 +122,9 @@ pub(crate) fn normalize_imports<'a>(
                         for comment in inline {
                             import_from.comments.inline.push(comment.value);
                         }
+
+                        import_from.trailing_comments =
+                            after.iter().map(|comment| comment.value.clone()).collect();
                     }
                 }
 

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

@@ -36,13 +36,14 @@ pub(crate) fn order_imports<'a>(
                     ImportFromStatement {
                         comments,
                         aliases,
-                        trailing_comma,
+                        trailing_comma, trailing_comments,
                     },
                 )| {
                     // Within each `Stmt::ImportFrom`, sort the members.
                     (
                         import_from,
                         comments,
+                        trailing_comments,
                         trailing_comma,
                         aliases
                             .into_iter()
@@ -56,7 +57,7 @@ pub(crate) fn order_imports<'a>(
 
     let ordered_imports = if matches!(section, ImportSection::Known(ImportType::Future)) {
         from_imports
-            .sorted_by_cached_key(|(import_from, _, _, aliases)| {
+            .sorted_by_cached_key(|(import_from, _, _, _, aliases)| {
                 ModuleKey::from_module(
                     import_from.module,
                     None,
@@ -95,7 +96,7 @@ pub(crate) fn order_imports<'a>(
                     ImportStyle::Straight,
                     settings,
                 ),
-                ImportFrom((import_from, _, _, aliases)) => ModuleKey::from_module(
+                ImportFrom((import_from, _, _, _, aliases)) => ModuleKey::from_module(
                     import_from.module,
                     None,
                     import_from.level,
@@ -117,7 +118,7 @@ pub(crate) fn order_imports<'a>(
             )
         });
         let ordered_from_imports =
-            from_imports.sorted_by_cached_key(|(import_from, _, _, aliases)| {
+            from_imports.sorted_by_cached_key(|(import_from, _, _, _, aliases)| {
                 ModuleKey::from_module(
                     import_from.module,
                     None,

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

@@ -0,0 +1,30 @@
+---
+source: crates/ruff_linter/src/rules/isort/mod.rs
+---
+trailing_comment.py:8:1: I001 [*] Import block is un-sorted or un-formatted
+   |
+ 6 |   pass
+ 7 |   
+ 8 | / from mylib import (
+ 9 | |     MyClient,
+10 | |     MyMgmtClient,
+11 | | ); from foo import (
+12 | |     bar
+13 | | )# some comment
+   |
+   = help: Organize imports
+
+ℹ Safe fix
+5  5  | 
+6  6  | pass
+7  7  | 
+   8  |+from foo import (
+   9  |+    bar,
+   10 |+)  # some comment
+8  11 | from mylib import (
+9  12 |     MyClient,
+10 13 |     MyMgmtClient,
+11    |-); from foo import (
+12    |-    bar
+13    |-)# some comment
+   14 |+)

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

@@ -68,6 +68,7 @@ pub(crate) struct ImportFromStatement<'a> {
     pub(crate) comments: CommentSet<'a>,
     pub(crate) aliases: FxHashMap<AliasData<'a>, CommentSet<'a>>,
     pub(crate) trailing_comma: TrailingComma,
+    pub(crate) trailing_comments: Vec<Cow<'a, str>>,
 }
 
 #[derive(Debug, Default)]
@@ -94,6 +95,7 @@ type Import<'a> = AliasDataWithComments<'a>;
 type ImportFrom<'a> = (
     ImportFromData<'a>,
     CommentSet<'a>,
+    Vec<Cow<'a, str>>,
     TrailingComma,
     Vec<AliasDataWithComments<'a>>,
 );