[pycodestyle] Add an autofix for E402

[PeterJCLaw]

Summary

This implements a fixer for E402, as suggested at #6514.

This enables users to move imports up to the top of the file when they are found lower down. This is similar to isort's float-to-top logic, though not quite the same. As suggested by @ntBre on that issue, the expected path here is for people to enable unsafe-fixes to access this feature, potentially marking this fix "safe" in their config, rather than to implement an float-to-top option in the isort config.

Internally this means that there are separate fix edits for each import found to be out of place rather than a single one (the latter being how isort tends to work).

The current state of this is a first draft. It's working, I've done some manual testing and I have copied over a lot of the handling from the isort-like organize_imports function, though I may have missed things.

I'm also fairly new to rust and to this project so e.g: I don't know if e.g: the code is the best way to spell some of the concepts. I'm very happy to refactor things if there's better ways, including if e.g: we want to share more code between this and the isort crate -- I've deliberately not done that so far (purely to get to something which works).

This remains a bit work-in-progress, so I've included a test project I'm using for testing against though obviously that would be removed before this gets anywhere near merging. (I might also tidy the history a bit too)

Fixes #6514

Test Plan

Manual testing so far, see various demo files currently committed. I'd like to add some tests and will look into that next. It looks like it's mostly snapshot based testing. Do we just add a new test function in crates/ruff_linter/src/rules/pycodestyle/mod.rs?

Noting them down so they're written, test cases:

• interaction with editor organise-imports logic -- with and without the E402 fix being accepted as safe
• variations on import spelling
• variations of other statements near the import, including on the same line (e.g: peruse the isort tests and pick relevant looking ones)

[astral-sh-bot]

ruff-ecosystem results

Linter (stable)

✅ ecosystem check detected no linter changes.

Linter (preview)

✅ ecosystem check detected no linter changes.

[MichaReiser]

Thanks for working on this. I'll close this PR due to inactivity. Feel free to resubmit it if you plan to continue working on this.

[ntBre]

Oops, I think this one was waiting on my review. Sorry @PeterJCLaw for the delay! If you'd still like me to have a look, please let me know, and I will reopen and add this back to my review queue.

[PeterJCLaw]

Apologies for the radio silence here. I am definitely hoping to keep working on this, I haven't had a lot of time recently. I may have some time next week or the following weekend, but not sure.

@ntBre if you have capacity for a review of what's here or any suggestions on approaching the missing pieces that'd be most welcome.

[ntBre]

I intended to have a look, but just got a bit behind. I'll reopen for now. Thank you!

[PeterJCLaw]

I've rebased this on the current main, fixed a missing part of the fixer implementation (fix_title method) and updated the snapshot tests.

The tests are mostly showing what I'd expected, however I think there may be some discrepancies in what ruff finds for E402 that flake8/pycodestyle finds differently. I'll put some details on that below, though I don't intend to change it as part of this PR.

Given the snapshot tests, I'm fairly happy that this is covered on that front. I'm not sure what the coverage expectations are though, so happy to look at adding more cases if that's desired.

The area I feel less able to immediately add a tests for, but which I think could benefit from having an integration test for, is the inclusion of this fixer in organize_imports_edit. That does seem to behave as we'd want, though I think it could be valuable to have a test to capture/demonstrate that behaviour.

E402 differences between tools

$ flake8 --show-source --select E402 crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_0.py
crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_0.py:25:1: E402 module level import not at top of file
import f
^
crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_0.py:27:1: E402 module level import not at top of file
import matplotlib
^
crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_0.py:31:1: E402 module level import not at top of file
import g
^
crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_0.py:35:1: E402 module level import not at top of file
import h
^
crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_0.py:45:1: E402 module level import not at top of file
import k; import l
^

However ruff produces:

$ cargo run -p ruff -- check --select E402 --no-cache crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_0.py
    Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.25s
     Running `target/debug/ruff check --select E402 --no-cache crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_0.py`
E402 Module level import not at top of file
  --> crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_0.py:35:1
   |
33 | __some__magic = 1
34 |
35 | import h
   | ^^^^^^^^
   |
help: Move module level imports to top of file

E402 Module level import not at top of file
  --> crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_0.py:45:1
   |
43 |     import j
44 |
45 | import k; import l
   | ^^^^^^^^
   |
help: Move module level imports to top of file

E402 Module level import not at top of file
  --> crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_0.py:45:11
   |
43 |     import j
44 |
45 | import k; import l
   |           ^^^^^^^^
   |
help: Move module level imports to top of file

Found 3 errors.
No fixes available (1 hidden fix can be enabled with the `--unsafe-fixes` option).

Interestingly, isort moves even more imports to the top.

isort --float-to-top --diff crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_0.py

--- .../ruff/crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_0.py:before
+++ .../ruff/crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_0.py:after
@@ -1,9 +1,19 @@
 """Top-level docstring."""
+import sys
+
+import a
+import c
+import f
+import g
+import h
+import k
+import l
+import matplotlib
+import n
 
 __all__ = ["y"]
 __version__: str = "0.1.0"
 
-import a
 
 try:
     import b
@@ -12,27 +22,21 @@
 else:
     pass
 
-import c
 
 if x > 0:
     import d
 else:
     import e
 
-import sys
 sys.path.insert(0, "some/path")
 
-import f
 
-import matplotlib
 
 matplotlib.use("Agg")
 
-import g
 
 __some__magic = 1
 
-import h
 
 
 def foo() -> None:
@@ -42,9 +46,7 @@
 if __name__ == "__main__":
     import j
 
-import k; import l
 
 
 if __name__ == "__main__":
     import m; \
-import n

I'm not sure whether this is deliberate or if this should be addressed separately. (I don't propose to change the detection in this PR).

[PeterJCLaw]

I've moved the fixer into being preview-only, I believe, and updated the tests.

Thanks for your help getting this far @ntBre; assuming that CI passes (tests pass locally) it would be great to get a review here when you have time.

[PeterJCLaw]

@ntBre just wondering when you might have a chance to review this? I believe it's ready -- the fixer is preview only and there are tests for both the preview and non-preview cases.

[ntBre]

Thank you for all of your work here and sorry again for the delay!

I only had a few minor comments and a suggestion to drop notebook support for now to simplify some things. I noticed a couple of other issues that I think are probably okay for now, but I'm curious to get your thoughts.

[ntBre]

I was a bit worried about this initially, but if your comment is correct, it seems okay. Just double-checking, did you verify that this only works in an editor if extend-safe-fixes is set? It looks right to me from a quick skim, but it seems easy to verify empirically.

[PeterJCLaw]

I believe I did early on, but I'll re-check and post back. Ideally I'd love to have a test for this but I couldn't easily find any tests which check this sort of thing. Do you have any suggestions for what such a test would look like (are there any examples I can borrow from or missed?)

[PeterJCLaw]

Verified using VScode with ruff built from 671fbf2 that the editor fix only applies when both preview = true and extend-safe-fixes = ["E402"] are set.
If there's a way add a test for this I'd be glad to do so.

[ntBre]

This API strikes me as a bit strange. Between this and the TODO comment, would it be helpful to skip the fix in notebooks for now? That would be totally fine with me if it would simplify the implementation.

[PeterJCLaw]

66a1fe8 improves the docstring. I've left the parameter since that's just a pass-through and I don't think should change anything in the callees.

[ntBre]

One potential issue here is that moving multiple imports reverses the order, if I'm following correctly. We ran into something similar recently with I002 (#20811), but the fix here wouldn't be so easy since we don't know all of the imports we're moving up front. I'm okay with this limitation because I don't see a good way around it, at least with the current structure of the rule, but we should probably document it somewhere.

[PeterJCLaw]

Ooh, good catch. I expect that a lot of people running this will also be using the isort rules, but given this is no longer coupled to those we can't rely on it.

I'm not immediately finding notes about linters, but in lots of cases there are notes about lints and their effects on the rule pages, so that seems a good place. I'll add something to https://docs.astral.sh/ruff/rules/module-import-not-at-top-of-file/. (I see you also suggesting adding something there about fix safety)

[PeterJCLaw]

Hopefully 6271fc9 covers this?

[ntBre]

It's slightly unfortunate to leave the comment hanging, but I don't think we can easily reuse the logic from the isort rules that tries to move comments too.

[PeterJCLaw]

Yeah, I've tried to copy over what I could. As far as I can tell the isort rules don't actually support it either (though isort itself does). I hope it's ok to leave this and let both be addressed together.

[ntBre]

I think you know about this since there's already a TODO comment, but it looks like this case is the one place where the import moves across cell boundaries. Again, I'm leaning toward excluding notebooks for now, just noting this for future reference.

[PeterJCLaw]

Looking at notebooks again, I'd actually hoped to get them working here. Though agree that it'd be simpler to drop them for now.

In an editor (VSCode at least) running the fix action in a notebook cell does work correctly currently, which suggests to me that there's something happening in the interaction with a file full of cells instead of whatever the editor interaction is. My first thought was that this was a cascading error since we're processing top to bottom, but that doesn't seem to fully explain what's happening here.

Anyway, 33b5efb disables the fixer for notebooks. (Though it looks like it still works in editor sessions 🤷)

[ntBre]

We should add a ## Fix safety section to the end of the docs noting that the fix is always unsafe because reordering imports can change runtime behavior. That's hopefully kind of obvious to readers but good to have the section anyway.

[PeterJCLaw]

6271fc9

[PeterJCLaw]

@ntBre gentle ping on review here, I believe I've responded to/addressed all the comments.

[ntBre]

Thank you! And sorry for the delay. This looks good to me now. My very last suggestion is around the importer API again. What do you think about something like this?

Patch

diff --git a/crates/ruff_linter/src/importer/mod.rs b/crates/ruff_linter/src/importer/mod.rs
index 95958c6640..3387a79a96 100644
--- a/crates/ruff_linter/src/importer/mod.rs
+++ b/crates/ruff_linter/src/importer/mod.rs
@@ -17,7 +17,8 @@ use ruff_python_parser::Parsed;
 use ruff_python_semantic::{
     ImportedName, MemberNameImport, ModuleNameImport, NameImport, SemanticModel,
 };
-use ruff_python_trivia::textwrap::indent;
+use ruff_python_trivia::{PythonWhitespace, textwrap::indent};
+use ruff_source_file::LineRanges;
 use ruff_text_size::{Ranged, TextRange, TextSize};
 
 use crate::cst::matchers::{match_aliases, match_import_from, match_statement};
@@ -77,22 +78,22 @@ impl<'a> Importer<'a> {
             // Insert after the last top-level import.
             Insertion::end_of_statement(stmt, self.source, self.stylist).into_edit(&required_import)
         } else {
-            self.add_at_start(&required_import, None)
+            self.add_at_start(&required_import)
         }
     }
 
-    /// Add a statement to the start of the file, potentially limited to a text range.
-    ///
-    /// The latter is useful for operating within notebooks.
-    pub(crate) fn add_at_start(&self, text: &str, within_range: Option<TextRange>) -> Edit {
-        // Check if there are any future imports that we need to respect
-        if let Some(last_future_import) = self.find_last_future_import(within_range) {
-            // Insert after the last future import
+    /// Add an existing import statement to the start of the file.
+    pub(crate) fn add_import_at_start(&self, import: &Stmt) -> Edit {
+        let range = TextRange::new(import.start(), self.source.line_end(import.end()));
+        self.add_at_start(self.source[range].trim_whitespace())
+    }
+
+    fn add_at_start(&self, text: &str) -> Edit {
+        if let Some(last_future_import) = self.find_last_future_import() {
             Insertion::end_of_statement(last_future_import, self.source, self.stylist)
                 .into_edit(text)
         } else {
-            // Insert at the start of the file.
-            Insertion::start_of_file(self.python_ast, self.source, self.stylist, within_range)
+            Insertion::start_of_file(self.python_ast, self.source, self.stylist, None)
                 .into_edit(text)
         }
     }
@@ -542,22 +543,8 @@ impl<'a> Importer<'a> {
     }
 
     /// Find the last `from __future__` import statement in the AST.
-    fn find_last_future_import(&self, within_range: Option<TextRange>) -> Option<&'a Stmt> {
-        // When operating within notebook cells we might only want to look at a
-        // range of the AST, yet can still encounter future imports at the top
-        // of a cell.
-        let full_body = self.python_ast;
-        let mut body = within_range
-            .map(|range| {
-                let start = full_body.partition_point(|stmt| stmt.start() < range.start());
-                let end = full_body.partition_point(|stmt| stmt.end() <= range.end());
-
-                &full_body[start..end]
-            })
-            .unwrap_or(full_body)
-            .iter()
-            .peekable();
-
+    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| {
diff --git a/crates/ruff_linter/src/rules/pycodestyle/rules/module_import_not_at_top_of_file.rs b/crates/ruff_linter/src/rules/pycodestyle/rules/module_import_not_at_top_of_file.rs
index 81886da596..0112808f32 100644
--- a/crates/ruff_linter/src/rules/pycodestyle/rules/module_import_not_at_top_of_file.rs
+++ b/crates/ruff_linter/src/rules/pycodestyle/rules/module_import_not_at_top_of_file.rs
@@ -1,6 +1,5 @@
 use ruff_macros::{ViolationMetadata, derive_message_formats};
 use ruff_python_ast::{PySourceType, Stmt};
-use ruff_python_trivia::PythonWhitespace;
 use ruff_source_file::LineRanges;
 use ruff_text_size::{Ranged, TextRange};
 
@@ -105,18 +104,10 @@ pub(crate) fn module_import_not_at_top_of_file(checker: &Checker, stmt: &Stmt) {
             return;
         }
 
-        let range = stmt.range();
+        let edit = checker.importer().add_import_at_start(stmt);
 
-        // Include comments but not the trailing newline (so we don't insert an extra newline).
-        let text_range = TextRange::new(range.start(), locator.line_end(range.end()));
-
-        let edit = checker
-            .importer()
-            .add_at_start(checker.source()[text_range].trim_whitespace(), None);
-
-        // Include comments *and* the trailing newline, so that we do remove the whole line.
-        let removal_range =
-            TextRange::new(text_range.start(), locator.full_line_end(text_range.end()));
+        // Include trailing comments and the newline in the removal.
+        let removal_range = TextRange::new(stmt.start(), locator.full_line_end(stmt.end()));
 
         diagnostic.set_fix(Fix::unsafe_edits(
             Edit::range_deletion(removal_range),

This reverts the now-unused within_range changes and makes the str version of add_at_start private, only exposing the add_import_at_start helper, which also simplifies the calling side a bit too.

[ntBre]

Thank you!

[PeterJCLaw]

Thanks for your help with this @ntBre :)