Skip to content

Commit

Permalink
Sort imports by cached key (#7963)
Browse files Browse the repository at this point in the history
## Summary

Refactor for isort implementation. Closes #7738.

I introduced a `NatOrdString` and a `NatOrdStr` type to have a naturally
ordered `String` and `&str`, and I pretty much went back to the original
implementation based on `module_key`, `member_key` and
`sorted_by_cached_key` from itertools. I tried my best to avoid
unnecessary allocations but it may have been clumsy in some places, so
feedback is appreciated! I also renamed the `Prefix` enum to
`MemberType` (and made some related adjustments) because I think this
fits more what it is, and it's closer to the wording found in the isort
documentation.

I think the result is nicer to work with, and it should make
implementing #1567 and the like easier :)

Of course, I am very much open to any and all remarks on what I did!

## Test Plan

I didn't add any test, I am relying on the existing tests since this is
just a refactor.
  • Loading branch information
bluthej committed Oct 30, 2023
1 parent 1f2d4f3 commit 5776ec1
Show file tree
Hide file tree
Showing 4 changed files with 152 additions and 165 deletions.
34 changes: 23 additions & 11 deletions crates/ruff_linter/src/rules/isort/mod.rs
Expand Up @@ -8,13 +8,14 @@ pub(crate) use categorize::categorize;
use categorize::categorize_imports;
pub use categorize::{ImportSection, ImportType};
use comments::Comment;
use itertools::Itertools;
use normalize::normalize_imports;
use order::order_imports;
use ruff_python_ast::PySourceType;
use ruff_python_codegen::Stylist;
use ruff_source_file::Locator;
use settings::Settings;
use sorting::cmp_either_import;
use sorting::module_key;
use types::EitherImport::{Import, ImportFrom};
use types::{AliasData, EitherImport, ImportBlock, TrailingComma};

Expand Down Expand Up @@ -173,17 +174,28 @@ fn format_import_block(

let imports = order_imports(import_block, settings);

let imports = {
let mut imports = imports
.import
.into_iter()
.map(Import)
.chain(imports.import_from.into_iter().map(ImportFrom))
.collect::<Vec<EitherImport>>();
if settings.force_sort_within_sections {
imports.sort_by(|import1, import2| cmp_either_import(import1, import2, settings));
};
let imports = imports
.import
.into_iter()
.map(Import)
.chain(imports.import_from.into_iter().map(ImportFrom));
let imports: Vec<EitherImport> = if settings.force_sort_within_sections {
imports
.sorted_by_cached_key(|import| match import {
Import((alias, _)) => {
module_key(Some(alias.name), alias.asname, None, None, settings)
}
ImportFrom((import_from, _, _, aliases)) => module_key(
import_from.module,
None,
import_from.level,
aliases.first().map(|(alias, _)| (alias.name, alias.asname)),
settings,
),
})
.collect()
} else {
imports.collect()
};

// Add a blank line between every section.
Expand Down
42 changes: 17 additions & 25 deletions crates/ruff_linter/src/rules/isort/order.rs
@@ -1,9 +1,7 @@
use std::cmp::Ordering;

use itertools::Itertools;

use super::settings::Settings;
use super::sorting::{cmp_import_from, cmp_members, cmp_modules};
use super::sorting::{member_key, module_key};
use super::types::{AliasData, CommentSet, ImportBlock, ImportFromStatement, OrderedImportBlock};

pub(crate) fn order_imports<'a>(
Expand All @@ -13,12 +11,11 @@ pub(crate) fn order_imports<'a>(
let mut ordered = OrderedImportBlock::default();

// Sort `Stmt::Import`.
ordered.import.extend(
block
.import
.into_iter()
.sorted_by(|(alias1, _), (alias2, _)| cmp_modules(alias1, alias2, settings)),
);
ordered
.import
.extend(block.import.into_iter().sorted_by_cached_key(|(alias, _)| {
module_key(Some(alias.name), alias.asname, None, None, settings)
}));

// Sort `Stmt::ImportFrom`.
ordered.import_from.extend(
Expand Down Expand Up @@ -53,27 +50,22 @@ pub(crate) fn order_imports<'a>(
trailing_comma,
aliases
.into_iter()
.sorted_by(|(alias1, _), (alias2, _)| {
cmp_members(alias1, alias2, settings)
.sorted_by_cached_key(|(alias, _)| {
member_key(alias.name, alias.asname, settings)
})
.collect::<Vec<(AliasData, CommentSet)>>(),
)
},
)
.sorted_by(
|(import_from1, _, _, aliases1), (import_from2, _, _, aliases2)| {
cmp_import_from(import_from1, import_from2, settings).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, settings)
}
}
})
},
),
.sorted_by_cached_key(|(import_from, _, _, aliases)| {
module_key(
import_from.module,
None,
import_from.level,
aliases.first().map(|(alias, _)| (alias.name, alias.asname)),
settings,
)
}),
);
ordered
}

0 comments on commit 5776ec1

Please sign in to comment.