diff --git a/crates/ruff_linter/resources/test/fixtures/isort/required_imports/docstring_future_import.py b/crates/ruff_linter/resources/test/fixtures/isort/required_imports/docstring_future_import.py new file mode 100644 index 00000000000000..fd1dea942b5809 --- /dev/null +++ b/crates/ruff_linter/resources/test/fixtures/isort/required_imports/docstring_future_import.py @@ -0,0 +1,3 @@ +"a docstring" +from __future__ import annotations +# EOF diff --git a/crates/ruff_linter/resources/test/fixtures/isort/required_imports/future_import.py b/crates/ruff_linter/resources/test/fixtures/isort/required_imports/future_import.py new file mode 100644 index 00000000000000..2e6e0107ef07cc --- /dev/null +++ b/crates/ruff_linter/resources/test/fixtures/isort/required_imports/future_import.py @@ -0,0 +1,2 @@ +from __future__ import annotations +# EOF diff --git a/crates/ruff_linter/src/importer/mod.rs b/crates/ruff_linter/src/importer/mod.rs index ec8745e0d7bb9e..4ab9f7af130b14 100644 --- a/crates/ruff_linter/src/importer/mod.rs +++ b/crates/ruff_linter/src/importer/mod.rs @@ -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( diff --git a/crates/ruff_linter/src/rules/isort/mod.rs b/crates/ruff_linter/src/rules/isort/mod.rs index 923a2d71d29d44..28c3c696507b05 100644 --- a/crates/ruff_linter/src/rules/isort/mod.rs +++ b/crates/ruff_linter/src/rules/isort/mod.rs @@ -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()); diff --git a/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__required_import_existing_import.py.snap b/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__required_import_existing_import.py.snap index 98a6e1ecbdbf8d..e4a5e566ff8dcd 100644 --- a/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__required_import_existing_import.py.snap +++ b/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__required_import_existing_import.py.snap @@ -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 diff --git a/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__required_import_with_alias_existing_import.py.snap b/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__required_import_with_alias_existing_import.py.snap index 1ff81a37621c6d..6eb136fab0120f 100644 --- a/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__required_import_with_alias_existing_import.py.snap +++ b/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__required_import_with_alias_existing_import.py.snap @@ -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 diff --git a/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__required_import_with_future_import_docstring_future_import.py.snap b/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__required_import_with_future_import_docstring_future_import.py.snap new file mode 100644 index 00000000000000..c78a254664cf61 --- /dev/null +++ b/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__required_import_with_future_import_docstring_future_import.py.snap @@ -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 diff --git a/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__required_import_with_future_import_future_import.py.snap b/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__required_import_with_future_import_future_import.py.snap new file mode 100644 index 00000000000000..af219a1219e7c0 --- /dev/null +++ b/crates/ruff_linter/src/rules/isort/snapshots/ruff_linter__rules__isort__tests__required_import_with_future_import_future_import.py.snap @@ -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