Normalize treatment of aliased and unaliased imports (#3216)

crates/ruff/resources/test/fixtures/isort/as_imports_comments.py

@@ -0,0 +1,15 @@
+from foo import (  # Comment on `foo`
+    Member as Alias,  # Comment on `Alias`
+)
+
+from bar import (  # Comment on `bar`
+    Member,  # Comment on `Member`
+)
+
+from baz import (  # Comment on `baz`
+    Member as Alias  # Comment on `Alias`
+)
+
+from bop import (  # Comment on `bop`
+    Member  # Comment on `Member`
+)

crates/ruff/src/rules/isort/categorize.rs

@@ -154,7 +154,7 @@ pub fn categorize_imports<'a>(
             .insert(import_from, aliases);
     }
     // Categorize `StmtKind::ImportFrom` (with re-export).
-    for ((import_from, alias), comments) in block.import_from_as {
+    for ((import_from, alias), aliases) in block.import_from_as {
         let classification = categorize(
             &import_from.module_base(),
             import_from.level,
@@ -170,7 +170,7 @@ pub fn categorize_imports<'a>(
             .entry(classification)
             .or_default()
             .import_from_as
-            .insert((import_from, alias), comments);
+            .insert((import_from, alias), aliases);
     }
     // Categorize `StmtKind::ImportFrom` (with star).
     for (import_from, comments) in block.import_from_star {

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

@@ -358,6 +358,7 @@ mod tests {
     use crate::test::{test_path, test_resource_path};
 
     #[test_case(Path::new("add_newline_before_comments.py"))]
+    #[test_case(Path::new("as_imports_comments.py"))]
     #[test_case(Path::new("combine_as_imports.py"))]
     #[test_case(Path::new("combine_import_from.py"))]
     #[test_case(Path::new("comments.py"))]

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

@@ -1,5 +1,6 @@
-use super::types::{AliasData, ImportBlock, ImportFromData, TrailingComma};
+use super::types::{AliasData, ImportBlock, ImportFromData};
 use super::AnnotatedImport;
+use crate::rules::isort::types::TrailingComma;
 
 pub fn normalize_imports(imports: Vec<AnnotatedImport>, combine_as_imports: bool) -> ImportBlock {
     let mut block = ImportBlock::default();
@@ -12,18 +13,18 @@ pub fn normalize_imports(imports: Vec<AnnotatedImport>, combine_as_imports: bool
             } => {
                 // Associate the comments with the first alias (best effort).
                 if let Some(name) = names.first() {
-                    let entry = block
+                    let comment_set = block
                         .import
                         .entry(AliasData {
                             name: name.name,
                             asname: name.asname,
                         })
                         .or_default();
                     for comment in atop {
-                        entry.atop.push(comment.value);
+                        comment_set.atop.push(comment.value);
                     }
                     for comment in inline {
-                        entry.inline.push(comment.value);
+                        comment_set.inline.push(comment.value);
                     }
                 }
 
@@ -46,18 +47,18 @@ pub fn normalize_imports(imports: Vec<AnnotatedImport>, combine_as_imports: bool
                 inline,
                 trailing_comma,
             } => {
+                // Insert comments on the statement itself.
                 if let Some(alias) = names.first() {
-                    let entry = if alias.name == "*" {
+                    let import_from = if alias.name == "*" {
                         block
                             .import_from_star
                             .entry(ImportFromData { module, level })
                             .or_default()
                     } else if alias.asname.is_none() || combine_as_imports {
-                        &mut block
+                        block
                             .import_from
                             .entry(ImportFromData { module, level })
                             .or_default()
-                            .0
                     } else {
                         block
                             .import_from_as
@@ -72,17 +73,17 @@ pub fn normalize_imports(imports: Vec<AnnotatedImport>, combine_as_imports: bool
                     };
 
                     for comment in atop {
-                        entry.atop.push(comment.value);
+                        import_from.comments.atop.push(comment.value);
                     }
 
                     for comment in inline {
-                        entry.inline.push(comment.value);
+                        import_from.comments.inline.push(comment.value);
                     }
                 }
 
-                // Create an entry for every alias.
+                // Create an entry for every alias (import) within the statement.
                 for alias in names {
-                    let entry = if alias.name == "*" {
+                    let import_from = if alias.name == "*" {
                         block
                             .import_from_star
                             .entry(ImportFromData { module, level })
@@ -92,12 +93,6 @@ pub fn normalize_imports(imports: Vec<AnnotatedImport>, combine_as_imports: bool
                             .import_from
                             .entry(ImportFromData { module, level })
                             .or_default()
-                            .1
-                            .entry(AliasData {
-                                name: alias.name,
-                                asname: alias.asname,
-                            })
-                            .or_default()
                     } else {
                         block
                             .import_from_as
@@ -111,20 +106,24 @@ pub fn normalize_imports(imports: Vec<AnnotatedImport>, combine_as_imports: bool
                             .or_default()
                     };
 
+                    let comment_set = import_from
+                        .aliases
+                        .entry(AliasData {
+                            name: alias.name,
+                            asname: alias.asname,
+                        })
+                        .or_default();
+
                     for comment in alias.atop {
-                        entry.atop.push(comment.value);
+                        comment_set.atop.push(comment.value);
                     }
                     for comment in alias.inline {
-                        entry.inline.push(comment.value);
+                        comment_set.inline.push(comment.value);
                     }
-                }
 
-                // Propagate trailing commas.
-                if matches!(trailing_comma, TrailingComma::Present) {
-                    if let Some(entry) =
-                        block.import_from.get_mut(&ImportFromData { module, level })
-                    {
-                        entry.2 = trailing_comma;
+                    // Propagate trailing commas.
+                    if matches!(trailing_comma, TrailingComma::Present) {
+                        import_from.trailing_comma = TrailingComma::Present;
                     }
                 }
             }

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

@@ -1,12 +1,12 @@
 use std::cmp::Ordering;
 use std::collections::BTreeSet;
 
+use crate::rules::isort::types::ImportFromStatement;
 use itertools::Itertools;
-use rustc_hash::FxHashMap;
 
 use super::settings::RelativeImportsOrder;
 use super::sorting::{cmp_import_from, cmp_members, cmp_modules};
-use super::types::{AliasData, CommentSet, ImportBlock, OrderedImportBlock, TrailingComma};
+use super::types::{AliasData, CommentSet, ImportBlock, OrderedImportBlock};
 
 pub fn order_imports<'a>(
     block: ImportBlock<'a>,
@@ -38,76 +38,43 @@ pub fn order_imports<'a>(
                 block
                     .import_from_as
                     .into_iter()
-                    .map(|((import_from, alias), comments)| {
-                        (
-                            import_from,
-                            (
-                                CommentSet {
-                                    atop: comments.atop,
-                                    inline: vec![],
-                                },
-                                FxHashMap::from_iter([(
-                                    alias,
-                                    CommentSet {
-                                        atop: vec![],
-                                        inline: comments.inline,
-                                    },
-                                )]),
-                                TrailingComma::Absent,
-                            ),
-                        )
-                    }),
+                    .map(|((import_from, ..), body)| (import_from, body)),
             )
             .chain(
                 // Include all star imports.
-                block
-                    .import_from_star
-                    .into_iter()
-                    .map(|(import_from, comments)| {
-                        (
-                            import_from,
-                            (
-                                CommentSet {
-                                    atop: comments.atop,
-                                    inline: vec![],
-                                },
-                                FxHashMap::from_iter([(
-                                    AliasData {
-                                        name: "*",
-                                        asname: None,
-                                    },
-                                    CommentSet {
-                                        atop: vec![],
-                                        inline: comments.inline,
-                                    },
-                                )]),
-                                TrailingComma::Absent,
-                            ),
-                        )
-                    }),
+                block.import_from_star.into_iter(),
             )
-            .map(|(import_from, (comments, aliases, locations))| {
-                // Within each `StmtKind::ImportFrom`, sort the members.
-                (
+            .map(
+                |(
                     import_from,
-                    comments,
-                    locations,
-                    aliases
-                        .into_iter()
-                        .sorted_by(|(alias1, _), (alias2, _)| {
-                            cmp_members(
-                                alias1,
-                                alias2,
-                                order_by_type,
-                                classes,
-                                constants,
-                                variables,
-                                force_to_top,
-                            )
-                        })
-                        .collect::<Vec<(AliasData, CommentSet)>>(),
-                )
-            })
+                    ImportFromStatement {
+                        comments,
+                        aliases,
+                        trailing_comma,
+                    },
+                )| {
+                    // Within each `StmtKind::ImportFrom`, sort the members.
+                    (
+                        import_from,
+                        comments,
+                        trailing_comma,
+                        aliases
+                            .into_iter()
+                            .sorted_by(|(alias1, _), (alias2, _)| {
+                                cmp_members(
+                                    alias1,
+                                    alias2,
+                                    order_by_type,
+                                    classes,
+                                    constants,
+                                    variables,
+                                    force_to_top,
+                                )
+                            })
+                            .collect::<Vec<(AliasData, CommentSet)>>(),
+                    )
+                },
+            )
             .sorted_by(
                 |(import_from1, _, _, aliases1), (import_from2, _, _, aliases2)| {
                     cmp_import_from(

crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__as_imports_comments.py.snap

@@ -0,0 +1,22 @@
+---
+source: crates/ruff/src/rules/isort/mod.rs
+expression: diagnostics
+---
+- kind:
+    UnsortedImports: ~
+  location:
+    row: 1
+    column: 0
+  end_location:
+    row: 16
+    column: 0
+  fix:
+    content: "from bar import (  # Comment on `bar`\n    Member,  # Comment on `Member`\n)\nfrom baz import Member as Alias  # Comment on `Alias`  # Comment on `baz`\nfrom bop import Member  # Comment on `Member`  # Comment on `bop`\nfrom foo import (  # Comment on `foo`\n    Member as Alias,  # Comment on `Alias`\n)\n"
+    location:
+      row: 1
+      column: 0
+    end_location:
+      row: 16
+      column: 0
+  parent: ~
+

crates/ruff/src/rules/isort/split.rs

@@ -46,10 +46,10 @@ pub fn split_by_forced_separate<'a>(
             .import_from
             .insert(imp, val);
     }
-    for ((imp, alias), comment_set) in import_from_as {
+    for ((imp, alias), val) in import_from_as {
         blocks[find_block_index(forced_separate, &imp)]
             .import_from_as
-            .insert((imp, alias), comment_set);
+            .insert((imp, alias), val);
     }
     for (imp, comment_set) in import_from_star {
         blocks[find_block_index(forced_separate, &imp)]

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

@@ -54,27 +54,27 @@ impl Importable for ImportFromData<'_> {
     }
 }
 
+#[derive(Debug, Default)]
+pub struct ImportFromStatement<'a> {
+    pub comments: CommentSet<'a>,
+    pub aliases: FxHashMap<AliasData<'a>, CommentSet<'a>>,
+    pub trailing_comma: TrailingComma,
+}
+
 #[derive(Debug, Default)]
 pub struct ImportBlock<'a> {
     // Set of (name, asname), used to track regular imports.
     // Ex) `import module`
     pub import: FxHashMap<AliasData<'a>, CommentSet<'a>>,
     // Map from (module, level) to `AliasData`, used to track 'from' imports.
     // Ex) `from module import member`
-    pub import_from: FxHashMap<
-        ImportFromData<'a>,
-        (
-            CommentSet<'a>,
-            FxHashMap<AliasData<'a>, CommentSet<'a>>,
-            TrailingComma,
-        ),
-    >,
+    pub import_from: FxHashMap<ImportFromData<'a>, ImportFromStatement<'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: FxHashMap<(ImportFromData<'a>, AliasData<'a>), CommentSet<'a>>,
+    pub import_from_as: FxHashMap<(ImportFromData<'a>, AliasData<'a>), ImportFromStatement<'a>>,
     // Map from (module, level) to `AliasData`, used to track star imports.
     // Ex) `from module import *`
-    pub import_from_star: FxHashMap<ImportFromData<'a>, CommentSet<'a>>,
+    pub import_from_star: FxHashMap<ImportFromData<'a>, ImportFromStatement<'a>>,
 }
 
 type AliasDataWithComments<'a> = (AliasData<'a>, CommentSet<'a>);