diff --git a/src/rules/isort/annotate.rs b/src/rules/isort/annotate.rs new file mode 100644 index 0000000000..43f1ba0331 --- /dev/null +++ b/src/rules/isort/annotate.rs @@ -0,0 +1,122 @@ +use rustpython_ast::{Stmt, StmtKind}; + +use super::comments::Comment; +use super::helpers::trailing_comma; +use super::types::{AliasData, TrailingComma}; +use super::{AnnotatedAliasData, AnnotatedImport}; +use crate::source_code::Locator; + +pub fn annotate_imports<'a>( + imports: &'a [&'a Stmt], + comments: Vec>, + locator: &Locator, + split_on_trailing_comma: bool, +) -> Vec> { + let mut annotated = vec![]; + let mut comments_iter = comments.into_iter().peekable(); + for import in imports { + match &import.node { + StmtKind::Import { names } => { + // Find comments above. + let mut atop = vec![]; + while let Some(comment) = + comments_iter.next_if(|comment| comment.location.row() < import.location.row()) + { + atop.push(comment); + } + + // Find comments inline. + let mut inline = vec![]; + while let Some(comment) = comments_iter.next_if(|comment| { + comment.end_location.row() == import.end_location.unwrap().row() + }) { + inline.push(comment); + } + + annotated.push(AnnotatedImport::Import { + names: names + .iter() + .map(|alias| AliasData { + name: &alias.node.name, + asname: alias.node.asname.as_deref(), + }) + .collect(), + atop, + inline, + }); + } + StmtKind::ImportFrom { + module, + names, + level, + } => { + // Find comments above. + let mut atop = vec![]; + while let Some(comment) = + comments_iter.next_if(|comment| comment.location.row() < import.location.row()) + { + atop.push(comment); + } + + // Find comments inline. + // We associate inline comments with the import statement unless there's a + // single member, and it's a single-line import (like `from foo + // import bar # noqa`). + let mut inline = vec![]; + if names.len() > 1 + || names + .first() + .map_or(false, |alias| alias.location.row() > import.location.row()) + { + while let Some(comment) = comments_iter + .next_if(|comment| comment.location.row() == import.location.row()) + { + inline.push(comment); + } + } + + // Capture names. + let mut aliases = vec![]; + for alias in names { + // Find comments above. + let mut alias_atop = vec![]; + while let Some(comment) = comments_iter + .next_if(|comment| comment.location.row() < alias.location.row()) + { + alias_atop.push(comment); + } + + // Find comments inline. + let mut alias_inline = vec![]; + while let Some(comment) = comments_iter.next_if(|comment| { + comment.end_location.row() == alias.end_location.unwrap().row() + }) { + alias_inline.push(comment); + } + + aliases.push(AnnotatedAliasData { + name: &alias.node.name, + asname: alias.node.asname.as_deref(), + atop: alias_atop, + inline: alias_inline, + }); + } + + annotated.push(AnnotatedImport::ImportFrom { + module: module.as_deref(), + names: aliases, + level: level.as_ref(), + trailing_comma: if split_on_trailing_comma { + trailing_comma(import, locator) + } else { + TrailingComma::default() + }, + atop, + inline, + }); + } + _ => unreachable!("Expected StmtKind::Import | StmtKind::ImportFrom"), + } + } + annotated +} diff --git a/src/rules/isort/categorize.rs b/src/rules/isort/categorize.rs index 24b4e99567..2e42903efc 100644 --- a/src/rules/isort/categorize.rs +++ b/src/rules/isort/categorize.rs @@ -1,4 +1,4 @@ -use std::collections::BTreeSet; +use std::collections::{BTreeMap, BTreeSet}; use std::fs; use std::path::{Path, PathBuf}; @@ -6,6 +6,7 @@ use log::debug; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; +use super::types::{ImportBlock, Importable}; use crate::python::sys::KNOWN_STANDARD_LIBRARY; #[derive( @@ -89,3 +90,83 @@ fn match_sources<'a>(paths: &'a [PathBuf], base: &str) -> Option<&'a Path> { } None } + +pub fn categorize_imports<'a>( + block: ImportBlock<'a>, + src: &[PathBuf], + package: Option<&Path>, + known_first_party: &BTreeSet, + known_third_party: &BTreeSet, + extra_standard_library: &BTreeSet, +) -> BTreeMap> { + let mut block_by_type: BTreeMap = BTreeMap::default(); + // Categorize `StmtKind::Import`. + for (alias, comments) in block.import { + let import_type = categorize( + &alias.module_base(), + None, + src, + package, + known_first_party, + known_third_party, + extra_standard_library, + ); + block_by_type + .entry(import_type) + .or_default() + .import + .insert(alias, comments); + } + // Categorize `StmtKind::ImportFrom` (without re-export). + for (import_from, aliases) in block.import_from { + let classification = categorize( + &import_from.module_base(), + import_from.level, + src, + package, + known_first_party, + known_third_party, + extra_standard_library, + ); + block_by_type + .entry(classification) + .or_default() + .import_from + .insert(import_from, aliases); + } + // Categorize `StmtKind::ImportFrom` (with re-export). + for ((import_from, alias), comments) in block.import_from_as { + let classification = categorize( + &import_from.module_base(), + import_from.level, + src, + package, + known_first_party, + known_third_party, + extra_standard_library, + ); + block_by_type + .entry(classification) + .or_default() + .import_from_as + .insert((import_from, alias), comments); + } + // Categorize `StmtKind::ImportFrom` (with star). + for (import_from, comments) in block.import_from_star { + let classification = categorize( + &import_from.module_base(), + import_from.level, + src, + package, + known_first_party, + known_third_party, + extra_standard_library, + ); + block_by_type + .entry(classification) + .or_default() + .import_from_star + .insert(import_from, comments); + } + block_by_type +} diff --git a/src/rules/isort/mod.rs b/src/rules/isort/mod.rs index cbdcc92558..0cc6174e66 100644 --- a/src/rules/isort/mod.rs +++ b/src/rules/isort/mod.rs @@ -1,30 +1,30 @@ //! Rules from [isort](https://pypi.org/project/isort/). -use std::cmp::Ordering; -use std::collections::{BTreeMap, BTreeSet}; + +use std::collections::BTreeSet; use std::path::{Path, PathBuf}; +use annotate::annotate_imports; +use categorize::categorize_imports; pub use categorize::{categorize, ImportType}; use comments::Comment; -use helpers::trailing_comma; use itertools::Either::{Left, Right}; -use itertools::Itertools; -use rustc_hash::FxHashMap; -use rustpython_ast::{Stmt, StmtKind}; +use normalize::normalize_imports; +use order::order_imports; use settings::RelativeImportsOrder; -use sorting::{cmp_either_import, cmp_import_from, cmp_members, cmp_modules}; +use sorting::cmp_either_import; use track::{Block, Trailer}; use types::EitherImport::{Import, ImportFrom}; -use types::{ - AliasData, CommentSet, EitherImport, ImportBlock, ImportFromData, Importable, - OrderedImportBlock, TrailingComma, -}; +use types::{AliasData, CommentSet, EitherImport, OrderedImportBlock, TrailingComma}; use crate::source_code::{Locator, Stylist}; +mod annotate; mod categorize; mod comments; mod format; mod helpers; +mod normalize; +mod order; pub(crate) mod rules; pub mod settings; mod sorting; @@ -56,454 +56,6 @@ pub enum AnnotatedImport<'a> { }, } -fn annotate_imports<'a>( - imports: &'a [&'a Stmt], - comments: Vec>, - locator: &Locator, - split_on_trailing_comma: bool, -) -> Vec> { - let mut annotated = vec![]; - let mut comments_iter = comments.into_iter().peekable(); - for import in imports { - match &import.node { - StmtKind::Import { names } => { - // Find comments above. - let mut atop = vec![]; - while let Some(comment) = - comments_iter.next_if(|comment| comment.location.row() < import.location.row()) - { - atop.push(comment); - } - - // Find comments inline. - let mut inline = vec![]; - while let Some(comment) = comments_iter.next_if(|comment| { - comment.end_location.row() == import.end_location.unwrap().row() - }) { - inline.push(comment); - } - - annotated.push(AnnotatedImport::Import { - names: names - .iter() - .map(|alias| AliasData { - name: &alias.node.name, - asname: alias.node.asname.as_deref(), - }) - .collect(), - atop, - inline, - }); - } - StmtKind::ImportFrom { - module, - names, - level, - } => { - // Find comments above. - let mut atop = vec![]; - while let Some(comment) = - comments_iter.next_if(|comment| comment.location.row() < import.location.row()) - { - atop.push(comment); - } - - // Find comments inline. - // We associate inline comments with the import statement unless there's a - // single member, and it's a single-line import (like `from foo - // import bar # noqa`). - let mut inline = vec![]; - if names.len() > 1 - || names - .first() - .map_or(false, |alias| alias.location.row() > import.location.row()) - { - while let Some(comment) = comments_iter - .next_if(|comment| comment.location.row() == import.location.row()) - { - inline.push(comment); - } - } - - // Capture names. - let mut aliases = vec![]; - for alias in names { - // Find comments above. - let mut alias_atop = vec![]; - while let Some(comment) = comments_iter - .next_if(|comment| comment.location.row() < alias.location.row()) - { - alias_atop.push(comment); - } - - // Find comments inline. - let mut alias_inline = vec![]; - while let Some(comment) = comments_iter.next_if(|comment| { - comment.end_location.row() == alias.end_location.unwrap().row() - }) { - alias_inline.push(comment); - } - - aliases.push(AnnotatedAliasData { - name: &alias.node.name, - asname: alias.node.asname.as_deref(), - atop: alias_atop, - inline: alias_inline, - }); - } - - annotated.push(AnnotatedImport::ImportFrom { - module: module.as_deref(), - names: aliases, - level: level.as_ref(), - trailing_comma: if split_on_trailing_comma { - trailing_comma(import, locator) - } else { - TrailingComma::default() - }, - atop, - inline, - }); - } - _ => unreachable!("Expected StmtKind::Import | StmtKind::ImportFrom"), - } - } - annotated -} - -fn normalize_imports(imports: Vec, combine_as_imports: bool) -> ImportBlock { - let mut block = ImportBlock::default(); - for import in imports { - match import { - AnnotatedImport::Import { - names, - atop, - inline, - } => { - // Associate the comments with the first alias (best effort). - if let Some(name) = names.first() { - let entry = block - .import - .entry(AliasData { - name: name.name, - asname: name.asname, - }) - .or_default(); - for comment in atop { - entry.atop.push(comment.value); - } - for comment in inline { - entry.inline.push(comment.value); - } - } - - // Create an entry for every alias. - for name in &names { - block - .import - .entry(AliasData { - name: name.name, - asname: name.asname, - }) - .or_default(); - } - } - AnnotatedImport::ImportFrom { - module, - names, - level, - atop, - inline, - trailing_comma, - } => { - if let Some(alias) = names.first() { - let entry = if alias.name == "*" { - block - .import_from_star - .entry(ImportFromData { module, level }) - .or_default() - } else if alias.asname.is_none() || combine_as_imports { - &mut block - .import_from - .entry(ImportFromData { module, level }) - .or_default() - .0 - } else { - block - .import_from_as - .entry(( - ImportFromData { module, level }, - AliasData { - name: alias.name, - asname: alias.asname, - }, - )) - .or_default() - }; - - for comment in atop { - entry.atop.push(comment.value); - } - - for comment in inline { - entry.inline.push(comment.value); - } - } - - // Create an entry for every alias. - for alias in names { - let entry = if alias.name == "*" { - block - .import_from_star - .entry(ImportFromData { module, level }) - .or_default() - } else if alias.asname.is_none() || combine_as_imports { - block - .import_from - .entry(ImportFromData { module, level }) - .or_default() - .1 - .entry(AliasData { - name: alias.name, - asname: alias.asname, - }) - .or_default() - } else { - block - .import_from_as - .entry(( - ImportFromData { module, level }, - AliasData { - name: alias.name, - asname: alias.asname, - }, - )) - .or_default() - }; - - for comment in alias.atop { - entry.atop.push(comment.value); - } - for comment in alias.inline { - entry.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; - } - } - } - } - } - block -} - -fn categorize_imports<'a>( - block: ImportBlock<'a>, - src: &[PathBuf], - package: Option<&Path>, - known_first_party: &BTreeSet, - known_third_party: &BTreeSet, - extra_standard_library: &BTreeSet, -) -> BTreeMap> { - let mut block_by_type: BTreeMap = BTreeMap::default(); - // Categorize `StmtKind::Import`. - for (alias, comments) in block.import { - let import_type = categorize( - &alias.module_base(), - None, - src, - package, - known_first_party, - known_third_party, - extra_standard_library, - ); - block_by_type - .entry(import_type) - .or_default() - .import - .insert(alias, comments); - } - // Categorize `StmtKind::ImportFrom` (without re-export). - for (import_from, aliases) in block.import_from { - let classification = categorize( - &import_from.module_base(), - import_from.level, - src, - package, - known_first_party, - known_third_party, - extra_standard_library, - ); - block_by_type - .entry(classification) - .or_default() - .import_from - .insert(import_from, aliases); - } - // Categorize `StmtKind::ImportFrom` (with re-export). - for ((import_from, alias), comments) in block.import_from_as { - let classification = categorize( - &import_from.module_base(), - import_from.level, - src, - package, - known_first_party, - known_third_party, - extra_standard_library, - ); - block_by_type - .entry(classification) - .or_default() - .import_from_as - .insert((import_from, alias), comments); - } - // Categorize `StmtKind::ImportFrom` (with star). - for (import_from, comments) in block.import_from_star { - let classification = categorize( - &import_from.module_base(), - import_from.level, - src, - package, - known_first_party, - known_third_party, - extra_standard_library, - ); - block_by_type - .entry(classification) - .or_default() - .import_from_star - .insert(import_from, comments); - } - block_by_type -} - -fn order_imports<'a>( - block: ImportBlock<'a>, - order_by_type: bool, - relative_imports_order: RelativeImportsOrder, - classes: &'a BTreeSet, - constants: &'a BTreeSet, - variables: &'a BTreeSet, -) -> OrderedImportBlock<'a> { - let mut ordered = OrderedImportBlock::default(); - - // Sort `StmtKind::Import`. - ordered.import.extend( - block - .import - .into_iter() - .sorted_by(|(alias1, _), (alias2, _)| cmp_modules(alias1, alias2)), - ); - - // Sort `StmtKind::ImportFrom`. - ordered.import_from.extend( - // Include all non-re-exports. - block - .import_from - .into_iter() - .chain( - // Include all re-exports. - 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, - ), - ) - }), - ) - .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, - ), - ) - }), - ) - .map(|(import_from, (comments, aliases, locations))| { - // Within each `StmtKind::ImportFrom`, sort the members. - ( - import_from, - comments, - locations, - aliases - .into_iter() - .sorted_by(|(alias1, _), (alias2, _)| { - cmp_members( - alias1, - alias2, - order_by_type, - classes, - constants, - variables, - ) - }) - .collect::>(), - ) - }) - .sorted_by( - |(import_from1, _, _, aliases1), (import_from2, _, _, aliases2)| { - cmp_import_from(import_from1, import_from2, relative_imports_order).then_with( - || match (aliases1.first(), aliases2.first()) { - (None, None) => Ordering::Equal, - (None, Some(_)) => Ordering::Less, - (Some(_), None) => Ordering::Greater, - (Some((alias1, _)), Some((alias2, _))) => cmp_members( - alias1, - alias2, - order_by_type, - classes, - constants, - variables, - ), - }, - ) - }, - ), - ); - ordered -} - fn force_single_line_imports<'a>( block: OrderedImportBlock<'a>, single_line_exclusions: &BTreeSet, diff --git a/src/rules/isort/normalize.rs b/src/rules/isort/normalize.rs new file mode 100644 index 0000000000..bcf1887391 --- /dev/null +++ b/src/rules/isort/normalize.rs @@ -0,0 +1,134 @@ +use super::types::{AliasData, ImportBlock, ImportFromData, TrailingComma}; +use super::AnnotatedImport; + +pub fn normalize_imports(imports: Vec, combine_as_imports: bool) -> ImportBlock { + let mut block = ImportBlock::default(); + for import in imports { + match import { + AnnotatedImport::Import { + names, + atop, + inline, + } => { + // Associate the comments with the first alias (best effort). + if let Some(name) = names.first() { + let entry = block + .import + .entry(AliasData { + name: name.name, + asname: name.asname, + }) + .or_default(); + for comment in atop { + entry.atop.push(comment.value); + } + for comment in inline { + entry.inline.push(comment.value); + } + } + + // Create an entry for every alias. + for name in &names { + block + .import + .entry(AliasData { + name: name.name, + asname: name.asname, + }) + .or_default(); + } + } + AnnotatedImport::ImportFrom { + module, + names, + level, + atop, + inline, + trailing_comma, + } => { + if let Some(alias) = names.first() { + let entry = if alias.name == "*" { + block + .import_from_star + .entry(ImportFromData { module, level }) + .or_default() + } else if alias.asname.is_none() || combine_as_imports { + &mut block + .import_from + .entry(ImportFromData { module, level }) + .or_default() + .0 + } else { + block + .import_from_as + .entry(( + ImportFromData { module, level }, + AliasData { + name: alias.name, + asname: alias.asname, + }, + )) + .or_default() + }; + + for comment in atop { + entry.atop.push(comment.value); + } + + for comment in inline { + entry.inline.push(comment.value); + } + } + + // Create an entry for every alias. + for alias in names { + let entry = if alias.name == "*" { + block + .import_from_star + .entry(ImportFromData { module, level }) + .or_default() + } else if alias.asname.is_none() || combine_as_imports { + block + .import_from + .entry(ImportFromData { module, level }) + .or_default() + .1 + .entry(AliasData { + name: alias.name, + asname: alias.asname, + }) + .or_default() + } else { + block + .import_from_as + .entry(( + ImportFromData { module, level }, + AliasData { + name: alias.name, + asname: alias.asname, + }, + )) + .or_default() + }; + + for comment in alias.atop { + entry.atop.push(comment.value); + } + for comment in alias.inline { + entry.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; + } + } + } + } + } + block +} diff --git a/src/rules/isort/order.rs b/src/rules/isort/order.rs new file mode 100644 index 0000000000..037f60586b --- /dev/null +++ b/src/rules/isort/order.rs @@ -0,0 +1,130 @@ +use std::cmp::Ordering; +use std::collections::BTreeSet; + +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}; + +pub fn order_imports<'a>( + block: ImportBlock<'a>, + order_by_type: bool, + relative_imports_order: RelativeImportsOrder, + classes: &'a BTreeSet, + constants: &'a BTreeSet, + variables: &'a BTreeSet, +) -> OrderedImportBlock<'a> { + let mut ordered = OrderedImportBlock::default(); + + // Sort `StmtKind::Import`. + ordered.import.extend( + block + .import + .into_iter() + .sorted_by(|(alias1, _), (alias2, _)| cmp_modules(alias1, alias2)), + ); + + // Sort `StmtKind::ImportFrom`. + ordered.import_from.extend( + // Include all non-re-exports. + block + .import_from + .into_iter() + .chain( + // Include all re-exports. + 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, + ), + ) + }), + ) + .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, + ), + ) + }), + ) + .map(|(import_from, (comments, aliases, locations))| { + // Within each `StmtKind::ImportFrom`, sort the members. + ( + import_from, + comments, + locations, + aliases + .into_iter() + .sorted_by(|(alias1, _), (alias2, _)| { + cmp_members( + alias1, + alias2, + order_by_type, + classes, + constants, + variables, + ) + }) + .collect::>(), + ) + }) + .sorted_by( + |(import_from1, _, _, aliases1), (import_from2, _, _, aliases2)| { + cmp_import_from(import_from1, import_from2, relative_imports_order).then_with( + || match (aliases1.first(), aliases2.first()) { + (None, None) => Ordering::Equal, + (None, Some(_)) => Ordering::Less, + (Some(_), None) => Ordering::Greater, + (Some((alias1, _)), Some((alias2, _))) => cmp_members( + alias1, + alias2, + order_by_type, + classes, + constants, + variables, + ), + }, + ) + }, + ), + ); + ordered +}