From 5f8294aea4d605e366d1ea1be6fec0a08da86302 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Mon, 5 Dec 2022 11:48:38 -0500 Subject: [PATCH] Preserve star imports when re-formatting import blocks (#1066) --- .../fixtures/isort/preserve_import_star.py | 6 ++ src/isort/mod.rs | 70 ++++++++++++++++++- ...isort__tests__preserve_import_star.py.snap | 20 ++++++ src/isort/types.rs | 3 + 4 files changed, 97 insertions(+), 2 deletions(-) create mode 100644 resources/test/fixtures/isort/preserve_import_star.py create mode 100644 src/isort/snapshots/ruff__isort__tests__preserve_import_star.py.snap diff --git a/resources/test/fixtures/isort/preserve_import_star.py b/resources/test/fixtures/isort/preserve_import_star.py new file mode 100644 index 0000000000..6810817249 --- /dev/null +++ b/resources/test/fixtures/isort/preserve_import_star.py @@ -0,0 +1,6 @@ +from some_other_module import some_class +from some_other_module import * +# Above +from some_module import some_class # Aside +# Above +from some_module import * # Aside diff --git a/src/isort/mod.rs b/src/isort/mod.rs index 32e0f9e904..e6892e8b77 100644 --- a/src/isort/mod.rs +++ b/src/isort/mod.rs @@ -191,7 +191,18 @@ fn normalize_imports(imports: Vec, combine_as_imports: bool) -> } => { // Associate the comments with the first alias (best effort). if let Some(alias) = names.first() { - if alias.asname.is_none() || combine_as_imports { + if alias.name == "*" { + let entry = block + .import_from_star + .entry(ImportFromData { module, level }) + .or_default(); + for comment in atop { + entry.atop.push(comment.value); + } + for comment in inline { + entry.inline.push(comment.value); + } + } else if alias.asname.is_none() || combine_as_imports { let entry = &mut block .import_from .entry(ImportFromData { module, level }) @@ -225,7 +236,18 @@ fn normalize_imports(imports: Vec, combine_as_imports: bool) -> // Create an entry for every alias. for alias in names { - if alias.asname.is_none() || combine_as_imports { + if alias.name == "*" { + let entry = block + .import_from_star + .entry(ImportFromData { module, level }) + .or_default(); + for comment in alias.atop { + entry.atop.push(comment.value); + } + for comment in alias.inline { + entry.inline.push(comment.value); + } + } else if alias.asname.is_none() || combine_as_imports { let entry = block .import_from .entry(ImportFromData { module, level }) @@ -323,6 +345,22 @@ fn categorize_imports<'a>( .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, + 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 } @@ -367,6 +405,33 @@ fn sort_imports(block: ImportBlock) -> OrderedImportBlock { ) }), ) + .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, + }, + )]), + ), + ) + }), + ) .map(|(import_from, (comments, aliases))| { // Within each `StmtKind::ImportFrom`, sort the members. ( @@ -486,6 +551,7 @@ mod tests { #[test_case(Path::new("order_by_type.py"))] #[test_case(Path::new("order_relative_imports_by_level.py"))] #[test_case(Path::new("preserve_comment_order.py"))] + #[test_case(Path::new("preserve_import_star.py"))] #[test_case(Path::new("preserve_indentation.py"))] #[test_case(Path::new("reorder_within_section.py"))] #[test_case(Path::new("separate_first_party_imports.py"))] diff --git a/src/isort/snapshots/ruff__isort__tests__preserve_import_star.py.snap b/src/isort/snapshots/ruff__isort__tests__preserve_import_star.py.snap new file mode 100644 index 0000000000..9c899227eb --- /dev/null +++ b/src/isort/snapshots/ruff__isort__tests__preserve_import_star.py.snap @@ -0,0 +1,20 @@ +--- +source: src/isort/mod.rs +expression: checks +--- +- kind: UnsortedImports + location: + row: 1 + column: 0 + end_location: + row: 7 + column: 0 + fix: + content: "# Above\nfrom some_module import * # Aside\n\n# Above\nfrom some_module import some_class # Aside\nfrom some_other_module import *\nfrom some_other_module import some_class\n" + location: + row: 1 + column: 0 + end_location: + row: 7 + column: 0 + diff --git a/src/isort/types.rs b/src/isort/types.rs index b11c57224f..651dfadb45 100644 --- a/src/isort/types.rs +++ b/src/isort/types.rs @@ -59,6 +59,9 @@ pub struct ImportBlock<'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>>, + // Map from (module, level) to `AliasData`, used to track star imports. + // Ex) `from module import *` + pub import_from_star: FxHashMap, CommentSet<'a>>, } type AliasDataWithComments<'a> = (AliasData<'a>, CommentSet<'a>);