diff --git a/crates/ruff_linter/resources/test/fixtures/isort/trailing_comment.py b/crates/ruff_linter/resources/test/fixtures/isort/trailing_comment.py new file mode 100644 index 00000000000000..72c45251d69864 --- /dev/null +++ b/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 + diff --git a/crates/ruff_linter/src/rules/isort/annotate.rs b/crates/ruff_linter/src/rules/isort/annotate.rs index a30cf78708547e..5a7b7bc1d1ba68 100644 --- a/crates/ruff_linter/src/rules/isort/annotate.rs +++ b/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"), diff --git a/crates/ruff_linter/src/rules/isort/format.rs b/crates/ruff_linter/src/rules/isort/format.rs index 1226bbc10317e9..9c0ae601928ebe 100644 --- a/crates/ruff_linter/src/rules/isort/format.rs +++ b/crates/ruff_linter/src/rules/isort/format.rs @@ -1,8 +1,10 @@ +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; @@ -10,7 +12,7 @@ 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,6 +145,22 @@ 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(' '); @@ -143,9 +168,16 @@ fn format_single_line( 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 diff --git a/crates/ruff_linter/src/rules/isort/mod.rs b/crates/ruff_linter/src/rules/isort/mod.rs index 4a82745e6cc928..43302d9075393b 100644 --- a/crates/ruff_linter/src/rules/isort/mod.rs +++ b/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>, pub(crate) inline: Vec>, + pub(crate) trailing: Vec>, } #[derive(Debug)] @@ -56,6 +57,7 @@ pub(crate) enum AnnotatedImport<'a> { level: u32, atop: Vec>, inline: Vec>, + trailing: Vec>, 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"))] diff --git a/crates/ruff_linter/src/rules/isort/normalize.rs b/crates/ruff_linter/src/rules/isort/normalize.rs index 01896580232bb9..f7f7bcabce9fba 100644 --- a/crates/ruff_linter/src/rules/isort/normalize.rs +++ b/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) { diff --git a/crates/ruff_linter/src/rules/isort/order.rs b/crates/ruff_linter/src/rules/isort/order.rs index 19f0dc29002322..6555101aad2519 100644 --- a/crates/ruff_linter/src/rules/isort/order.rs +++ b/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::>(), + .collect::>(), ) }, ); diff --git a/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__as_imports_comments.py.snap b/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__as_imports_comments.py.snap index da70289eb48d69..4866de74a682e0 100644 --- a/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__as_imports_comments.py.snap +++ b/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 | ) - - diff --git a/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__trailing_comment.py.snap b/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__trailing_comment.py.snap new file mode 100644 index 00000000000000..890aaa10dc5ebb --- /dev/null +++ b/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__trailing_comment.py.snap @@ -0,0 +1,151 @@ +--- +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 +14 | | +15 | | pass + | |_^ I001 +16 | +17 | from foo import ( + | + = help: Organize imports + +ℹ Safe fix +5 5 | +6 6 | pass +7 7 | + 8 |+from foo import bar # some comment +8 9 | from mylib import ( +9 10 | MyClient, +10 11 | MyMgmtClient, +11 |-); from foo import ( +12 |- bar +13 |-)# some comment + 12 |+) +14 13 | +15 14 | pass +16 15 | + +trailing_comment.py:17:1: I001 [*] Import block is un-sorted or un-formatted + | +15 | pass +16 | +17 | / from foo import ( +18 | | bar +19 | | ) +20 | | from mylib import ( +21 | | MyClient, +22 | | MyMgmtClient, +23 | | # some comment +24 | | ) +25 | | +26 | | pass + | |_^ I001 +27 | +28 | from mylib import ( + | + = help: Organize imports + +ℹ Safe fix +14 14 | +15 15 | pass +16 16 | +17 |-from foo import ( +18 |- bar +19 |-) + 17 |+from foo import bar +20 18 | from mylib import ( +21 19 | MyClient, +22 20 | MyMgmtClient, + +trailing_comment.py:35:1: I001 [*] Import block is un-sorted or un-formatted + | +33 | pass +34 | +35 | / from mylib import ( +36 | | MyClient +37 | | # some comment +38 | | ) +39 | | +40 | | pass + | |_^ I001 +41 | +42 | from mylib import ( + | + = help: Organize imports + +ℹ Safe fix +32 32 | +33 33 | pass +34 34 | +35 |-from mylib import ( +36 |- MyClient +37 |- # some comment +38 |-) + 35 |+from mylib import MyClient # some comment +39 36 | +40 37 | pass +41 38 | + +trailing_comment.py:42:1: I001 [*] Import block is un-sorted or un-formatted + | +40 | pass +41 | +42 | / from mylib import ( +43 | | # some comment +44 | | MyClient +45 | | ) +46 | | +47 | | pass + | |_^ I001 +48 | +49 | # a + | + = help: Organize imports + +ℹ Safe fix +39 39 | +40 40 | pass +41 41 | +42 |-from mylib import ( +43 |- # some comment +44 |- MyClient +45 |-) + 42 |+from mylib import MyClient # some comment +46 43 | +47 44 | pass +48 45 | + +trailing_comment.py:50:1: I001 [*] Import block is un-sorted or un-formatted + | +49 | # a +50 | / from mylib import ( # b +51 | | # c +52 | | MyClient # d +53 | | # e +54 | | ) # f +55 | | + | |_^ I001 + | + = help: Organize imports + +ℹ Safe fix +47 47 | pass +48 48 | +49 49 | # a +50 |-from mylib import ( # b +51 |- # c +52 |- MyClient # d +53 |- # e +54 |-) # f + 50 |+from mylib import MyClient # b # c # d # e # f +55 51 | diff --git a/crates/ruff_linter/src/rules/isort/types.rs b/crates/ruff_linter/src/rules/isort/types.rs index 9fc4cc7ff86bcd..17b1567bcd63f2 100644 --- a/crates/ruff_linter/src/rules/isort/types.rs +++ b/crates/ruff_linter/src/rules/isort/types.rs @@ -24,11 +24,18 @@ pub(crate) struct AliasData<'a> { } #[derive(Debug, Default, Clone)] -pub(crate) struct CommentSet<'a> { +pub(crate) struct ImportCommentSet<'a> { pub(crate) atop: Vec>, pub(crate) inline: Vec>, } +#[derive(Debug, Default, Clone)] +pub(crate) struct ImportFromCommentSet<'a> { + pub(crate) atop: Vec>, + pub(crate) inline: Vec>, + pub(crate) trailing: Vec>, +} + pub(crate) trait Importable<'a> { fn module_name(&self) -> Cow<'a, str>; @@ -65,8 +72,8 @@ impl<'a> Importable<'a> for ImportFromData<'a> { #[derive(Debug, Default)] pub(crate) struct ImportFromStatement<'a> { - pub(crate) comments: CommentSet<'a>, - pub(crate) aliases: FxHashMap, CommentSet<'a>>, + pub(crate) comments: ImportFromCommentSet<'a>, + pub(crate) aliases: FxHashMap, ImportFromCommentSet<'a>>, pub(crate) trailing_comma: TrailingComma, } @@ -74,7 +81,7 @@ pub(crate) struct ImportFromStatement<'a> { pub(crate) struct ImportBlock<'a> { // Set of (name, asname), used to track regular imports. // Ex) `import module` - pub(crate) import: FxHashMap, CommentSet<'a>>, + pub(crate) import: FxHashMap, ImportCommentSet<'a>>, // Map from (module, level) to `AliasData`, used to track 'from' imports. // Ex) `from module import member` pub(crate) import_from: FxHashMap, ImportFromStatement<'a>>, @@ -87,15 +94,13 @@ pub(crate) struct ImportBlock<'a> { pub(crate) import_from_star: FxHashMap, ImportFromStatement<'a>>, } -type AliasDataWithComments<'a> = (AliasData<'a>, CommentSet<'a>); - -type Import<'a> = AliasDataWithComments<'a>; +type Import<'a> = (AliasData<'a>, ImportCommentSet<'a>); type ImportFrom<'a> = ( ImportFromData<'a>, - CommentSet<'a>, + ImportFromCommentSet<'a>, TrailingComma, - Vec>, + Vec<(AliasData<'a>, ImportFromCommentSet<'a>)>, ); #[derive(Debug)]