Skip to content

Commit 5776ec1

Browse files
authored
Sort imports by cached key (#7963)
## 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.
1 parent 1f2d4f3 commit 5776ec1

File tree

4 files changed

+152
-165
lines changed

4 files changed

+152
-165
lines changed

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

Lines changed: 23 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -8,13 +8,14 @@ pub(crate) use categorize::categorize;
88
use categorize::categorize_imports;
99
pub use categorize::{ImportSection, ImportType};
1010
use comments::Comment;
11+
use itertools::Itertools;
1112
use normalize::normalize_imports;
1213
use order::order_imports;
1314
use ruff_python_ast::PySourceType;
1415
use ruff_python_codegen::Stylist;
1516
use ruff_source_file::Locator;
1617
use settings::Settings;
17-
use sorting::cmp_either_import;
18+
use sorting::module_key;
1819
use types::EitherImport::{Import, ImportFrom};
1920
use types::{AliasData, EitherImport, ImportBlock, TrailingComma};
2021

@@ -173,17 +174,28 @@ fn format_import_block(
173174

174175
let imports = order_imports(import_block, settings);
175176

176-
let imports = {
177-
let mut imports = imports
178-
.import
179-
.into_iter()
180-
.map(Import)
181-
.chain(imports.import_from.into_iter().map(ImportFrom))
182-
.collect::<Vec<EitherImport>>();
183-
if settings.force_sort_within_sections {
184-
imports.sort_by(|import1, import2| cmp_either_import(import1, import2, settings));
185-
};
177+
let imports = imports
178+
.import
179+
.into_iter()
180+
.map(Import)
181+
.chain(imports.import_from.into_iter().map(ImportFrom));
182+
let imports: Vec<EitherImport> = if settings.force_sort_within_sections {
186183
imports
184+
.sorted_by_cached_key(|import| match import {
185+
Import((alias, _)) => {
186+
module_key(Some(alias.name), alias.asname, None, None, settings)
187+
}
188+
ImportFrom((import_from, _, _, aliases)) => module_key(
189+
import_from.module,
190+
None,
191+
import_from.level,
192+
aliases.first().map(|(alias, _)| (alias.name, alias.asname)),
193+
settings,
194+
),
195+
})
196+
.collect()
197+
} else {
198+
imports.collect()
187199
};
188200

189201
// Add a blank line between every section.

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

Lines changed: 17 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
1-
use std::cmp::Ordering;
2-
31
use itertools::Itertools;
42

53
use super::settings::Settings;
6-
use super::sorting::{cmp_import_from, cmp_members, cmp_modules};
4+
use super::sorting::{member_key, module_key};
75
use super::types::{AliasData, CommentSet, ImportBlock, ImportFromStatement, OrderedImportBlock};
86

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

1513
// Sort `Stmt::Import`.
16-
ordered.import.extend(
17-
block
18-
.import
19-
.into_iter()
20-
.sorted_by(|(alias1, _), (alias2, _)| cmp_modules(alias1, alias2, settings)),
21-
);
14+
ordered
15+
.import
16+
.extend(block.import.into_iter().sorted_by_cached_key(|(alias, _)| {
17+
module_key(Some(alias.name), alias.asname, None, None, settings)
18+
}));
2219

2320
// Sort `Stmt::ImportFrom`.
2421
ordered.import_from.extend(
@@ -53,27 +50,22 @@ pub(crate) fn order_imports<'a>(
5350
trailing_comma,
5451
aliases
5552
.into_iter()
56-
.sorted_by(|(alias1, _), (alias2, _)| {
57-
cmp_members(alias1, alias2, settings)
53+
.sorted_by_cached_key(|(alias, _)| {
54+
member_key(alias.name, alias.asname, settings)
5855
})
5956
.collect::<Vec<(AliasData, CommentSet)>>(),
6057
)
6158
},
6259
)
63-
.sorted_by(
64-
|(import_from1, _, _, aliases1), (import_from2, _, _, aliases2)| {
65-
cmp_import_from(import_from1, import_from2, settings).then_with(|| {
66-
match (aliases1.first(), aliases2.first()) {
67-
(None, None) => Ordering::Equal,
68-
(None, Some(_)) => Ordering::Less,
69-
(Some(_), None) => Ordering::Greater,
70-
(Some((alias1, _)), Some((alias2, _))) => {
71-
cmp_members(alias1, alias2, settings)
72-
}
73-
}
74-
})
75-
},
76-
),
60+
.sorted_by_cached_key(|(import_from, _, _, aliases)| {
61+
module_key(
62+
import_from.module,
63+
None,
64+
import_from.level,
65+
aliases.first().map(|(alias, _)| (alias.name, alias.asname)),
66+
settings,
67+
)
68+
}),
7769
);
7870
ordered
7971
}

0 commit comments

Comments
 (0)