[`isort`] Fix inserting required imports before future imports (`I002`) (#20676)

## Summary

Fixes #20674

---------

Co-authored-by: Brent Westbrook <36778786+ntBre@users.noreply.github.com>
This commit is contained in:
Dan Parizher 2025-10-06 09:40:36 -04:00 committed by GitHub
parent 1c5666ce5d
commit 9a29f7a339
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 77 additions and 9 deletions

View File

@ -0,0 +1,3 @@
"a docstring"
from __future__ import annotations
# EOF

View File

@ -0,0 +1,2 @@
from __future__ import annotations
# EOF

View File

@ -67,17 +67,25 @@ impl<'a> Importer<'a> {
/// Add an import statement to import the given module.
///
/// If there are no existing imports, the new import will be added at the top
/// of the file. Otherwise, it will be added after the most recent top-level
/// import statement.
/// of the file. If there are future imports, the new import will be added
/// after the last future import. Otherwise, it will be added after the most
/// recent top-level import statement.
pub(crate) fn add_import(&self, import: &NameImport, at: TextSize) -> Edit {
let required_import = import.to_string();
if let Some(stmt) = self.preceding_import(at) {
// Insert after the last top-level import.
Insertion::end_of_statement(stmt, self.source, self.stylist).into_edit(&required_import)
} else {
// Insert at the start of the file.
Insertion::start_of_file(self.python_ast, self.source, self.stylist)
.into_edit(&required_import)
// Check if there are any future imports that we need to respect
if let Some(last_future_import) = self.find_last_future_import() {
// Insert after the last future import
Insertion::end_of_statement(last_future_import, self.source, self.stylist)
.into_edit(&required_import)
} else {
// Insert at the start of the file.
Insertion::start_of_file(self.python_ast, self.source, self.stylist)
.into_edit(&required_import)
}
}
}
@ -524,6 +532,18 @@ impl<'a> Importer<'a> {
}
}
/// Find the last `from __future__` import statement in the AST.
fn find_last_future_import(&self) -> Option<&'a Stmt> {
let mut body = self.python_ast.iter().peekable();
let _docstring = body.next_if(|stmt| ast::helpers::is_docstring_stmt(stmt));
body.take_while(|stmt| {
stmt.as_import_from_stmt()
.is_some_and(|import_from| import_from.module.as_deref() == Some("__future__"))
})
.last()
}
/// Add a `from __future__ import annotations` import.
pub(crate) fn add_future_import(&self) -> Edit {
let import = &NameImport::ImportFrom(MemberNameImport::member(

View File

@ -1014,6 +1014,30 @@ mod tests {
Ok(())
}
#[test_case(Path::new("future_import.py"))]
#[test_case(Path::new("docstring_future_import.py"))]
fn required_import_with_future_import(path: &Path) -> Result<()> {
let snapshot = format!(
"required_import_with_future_import_{}",
path.to_string_lossy()
);
let diagnostics = test_path(
Path::new("isort/required_imports").join(path).as_path(),
&LinterSettings {
src: vec![test_resource_path("fixtures/isort")],
isort: super::settings::Settings {
required_imports: BTreeSet::from_iter([NameImport::Import(
ModuleNameImport::module("this".to_string()),
)]),
..super::settings::Settings::default()
},
..LinterSettings::for_rule(Rule::MissingRequiredImport)
},
)?;
assert_diagnostics!(snapshot, diagnostics);
Ok(())
}
#[test_case(Path::new("from_first.py"))]
fn from_first(path: &Path) -> Result<()> {
let snapshot = format!("from_first_{}", path.to_string_lossy());

View File

@ -4,6 +4,6 @@ source: crates/ruff_linter/src/rules/isort/mod.rs
I002 [*] Missing required import: `from __future__ import annotations`
--> existing_import.py:1:1
help: Insert required import: `from __future__ import annotations`
1 + from __future__ import annotations
2 | from __future__ import generator_stop
1 | from __future__ import generator_stop
2 + from __future__ import annotations
3 | import os

View File

@ -4,6 +4,6 @@ source: crates/ruff_linter/src/rules/isort/mod.rs
I002 [*] Missing required import: `from __future__ import annotations as _annotations`
--> existing_import.py:1:1
help: Insert required import: `from __future__ import annotations as _annotations`
1 + from __future__ import annotations as _annotations
2 | from __future__ import generator_stop
1 | from __future__ import generator_stop
2 + from __future__ import annotations as _annotations
3 | import os

View File

@ -0,0 +1,10 @@
---
source: crates/ruff_linter/src/rules/isort/mod.rs
---
I002 [*] Missing required import: `import this`
--> docstring_future_import.py:1:1
help: Insert required import: `import this`
1 | "a docstring"
2 | from __future__ import annotations
3 + import this
4 | # EOF

View File

@ -0,0 +1,9 @@
---
source: crates/ruff_linter/src/rules/isort/mod.rs
---
I002 [*] Missing required import: `import this`
--> future_import.py:1:1
help: Insert required import: `import this`
1 | from __future__ import annotations
2 + import this
3 | # EOF