From 42f61535b56887a73c4f0646b3deff49affd456d Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Fri, 24 Feb 2023 17:11:03 -0500 Subject: [PATCH] Normalize treatment of aliased and unaliased imports (#3216) --- .../fixtures/isort/as_imports_comments.py | 15 +++ crates/ruff/src/rules/isort/categorize.rs | 4 +- crates/ruff/src/rules/isort/mod.rs | 1 + crates/ruff/src/rules/isort/normalize.rs | 51 +++++---- crates/ruff/src/rules/isort/order.rs | 101 ++++++------------ ..._isort__tests__as_imports_comments.py.snap | 22 ++++ crates/ruff/src/rules/isort/split.rs | 4 +- crates/ruff/src/rules/isort/types.rs | 20 ++-- 8 files changed, 111 insertions(+), 107 deletions(-) create mode 100644 crates/ruff/resources/test/fixtures/isort/as_imports_comments.py create mode 100644 crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__as_imports_comments.py.snap diff --git a/crates/ruff/resources/test/fixtures/isort/as_imports_comments.py b/crates/ruff/resources/test/fixtures/isort/as_imports_comments.py new file mode 100644 index 0000000000000..70fe11393efaf --- /dev/null +++ b/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` +) diff --git a/crates/ruff/src/rules/isort/categorize.rs b/crates/ruff/src/rules/isort/categorize.rs index 114f0439e0f3b..c25ff18e112a9 100644 --- a/crates/ruff/src/rules/isort/categorize.rs +++ b/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 { diff --git a/crates/ruff/src/rules/isort/mod.rs b/crates/ruff/src/rules/isort/mod.rs index 93ed32ccb6217..9003e8c602e6e 100644 --- a/crates/ruff/src/rules/isort/mod.rs +++ b/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"))] diff --git a/crates/ruff/src/rules/isort/normalize.rs b/crates/ruff/src/rules/isort/normalize.rs index bcf18873917fb..b78ee433a51f5 100644 --- a/crates/ruff/src/rules/isort/normalize.rs +++ b/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, combine_as_imports: bool) -> ImportBlock { let mut block = ImportBlock::default(); @@ -12,7 +13,7 @@ pub fn normalize_imports(imports: Vec, 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, @@ -20,10 +21,10 @@ pub fn normalize_imports(imports: Vec, combine_as_imports: bool }) .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, 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, 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, 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, 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; } } } diff --git a/crates/ruff/src/rules/isort/order.rs b/crates/ruff/src/rules/isort/order.rs index 9322ab6a62b79..01b67ed7ae0f6 100644 --- a/crates/ruff/src/rules/isort/order.rs +++ b/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::>(), - ) - }) + 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::>(), + ) + }, + ) .sorted_by( |(import_from1, _, _, aliases1), (import_from2, _, _, aliases2)| { cmp_import_from( diff --git a/crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__as_imports_comments.py.snap b/crates/ruff/src/rules/isort/snapshots/ruff__rules__isort__tests__as_imports_comments.py.snap new file mode 100644 index 0000000000000..c97e326ce726c --- /dev/null +++ b/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: ~ + diff --git a/crates/ruff/src/rules/isort/split.rs b/crates/ruff/src/rules/isort/split.rs index d1092bf56a9de..70fd474e34a48 100644 --- a/crates/ruff/src/rules/isort/split.rs +++ b/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)] diff --git a/crates/ruff/src/rules/isort/types.rs b/crates/ruff/src/rules/isort/types.rs index 4ed3ef40d4a03..f02fe2c173fcf 100644 --- a/crates/ruff/src/rules/isort/types.rs +++ b/crates/ruff/src/rules/isort/types.rs @@ -54,6 +54,13 @@ impl Importable for ImportFromData<'_> { } } +#[derive(Debug, Default)] +pub struct ImportFromStatement<'a> { + pub comments: CommentSet<'a>, + pub aliases: FxHashMap, CommentSet<'a>>, + pub trailing_comma: TrailingComma, +} + #[derive(Debug, Default)] pub struct ImportBlock<'a> { // Set of (name, asname), used to track regular imports. @@ -61,20 +68,13 @@ pub struct ImportBlock<'a> { pub import: FxHashMap, 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, CommentSet<'a>>, - TrailingComma, - ), - >, + pub import_from: FxHashMap, 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, CommentSet<'a>>, + pub import_from_star: FxHashMap, ImportFromStatement<'a>>, } type AliasDataWithComments<'a> = (AliasData<'a>, CommentSet<'a>);