[pycodestyle] Add fix for multiple-imports-on-one-line (E401) (#…

…9518)

## Summary

Add autofix for `multiple_imports_on_one_line`, `E401`

## Test Plan

`cargo test`

crates/ruff_linter/resources/test/fixtures/pycodestyle/E40.py

@@ -1,5 +1,6 @@
 #: E401
 import os, sys
+
 #: Okay
 import os
 import sys
@@ -59,3 +60,21 @@
 a = 1
 
 import bar
+
+#: E401
+import re as regex, string  # also with a comment!
+import re as regex, string; x = 1
+
+x = 1; import re as regex, string
+
+
+def blah():
+    import datetime as dt, copy
+
+    def nested_and_tested():
+        import builtins, textwrap as tw
+
+        x = 1; import re as regex, string
+        import re as regex, string; x = 1
+
+    if True: import re as regex, string

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

@@ -67,6 +67,7 @@ mod tests {
     }
 
     #[test_case(Rule::IsLiteral, Path::new("constant_literals.py"))]
+    #[test_case(Rule::MultipleImportsOnOneLine, Path::new("E40.py"))]
     #[test_case(Rule::ModuleImportNotAtTopOfFile, Path::new("E402_0.py"))]
     #[test_case(Rule::TypeComparison, Path::new("E721.py"))]
     fn preview_rules(rule_code: Rule, path: &Path) -> Result<()> {

crates/ruff_linter/src/rules/pycodestyle/rules/multiple_imports_on_one_line.rs

@@ -1,7 +1,13 @@
-use ruff_diagnostics::{Diagnostic, Violation};
+use itertools::Itertools;
+
+use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::{Alias, Stmt};
-use ruff_text_size::Ranged;
+use ruff_python_codegen::Stylist;
+use ruff_python_index::Indexer;
+use ruff_python_trivia::indentation_at_offset;
+use ruff_source_file::Locator;
+use ruff_text_size::{Ranged, TextRange};
 
 use crate::checkers::ast::Checker;
 
@@ -27,17 +33,88 @@ use crate::checkers::ast::Checker;
 pub struct MultipleImportsOnOneLine;
 
 impl Violation for MultipleImportsOnOneLine {
+    const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
+
     #[derive_message_formats]
     fn message(&self) -> String {
         format!("Multiple imports on one line")
     }
+
+    fn fix_title(&self) -> Option<String> {
+        Some(format!("Split imports"))
+    }
 }
 
 /// E401
 pub(crate) fn multiple_imports_on_one_line(checker: &mut Checker, stmt: &Stmt, names: &[Alias]) {
     if names.len() > 1 {
-        checker
-            .diagnostics
-            .push(Diagnostic::new(MultipleImportsOnOneLine, stmt.range()));
+        let mut diagnostic = Diagnostic::new(MultipleImportsOnOneLine, stmt.range());
+        if checker.settings.preview.is_enabled() {
+            diagnostic.set_fix(split_imports(
+                stmt,
+                names,
+                checker.locator(),
+                checker.indexer(),
+                checker.stylist(),
+            ));
+        }
+        checker.diagnostics.push(diagnostic);
+    }
+}
+
+/// Generate a [`Fix`] to split the imports across multiple statements.
+fn split_imports(
+    stmt: &Stmt,
+    names: &[Alias],
+    locator: &Locator,
+    indexer: &Indexer,
+    stylist: &Stylist,
+) -> Fix {
+    if indexer.in_multi_statement_line(stmt, locator) {
+        // Ex) `x = 1; import os, sys` (convert to `x = 1; import os; import sys`)
+        let replacement = names
+            .iter()
+            .map(|alias| {
+                let Alias {
+                    range: _,
+                    name,
+                    asname,
+                } = alias;
+
+                if let Some(asname) = asname {
+                    format!("import {name} as {asname}")
+                } else {
+                    format!("import {name}")
+                }
+            })
+            .join("; ");
+
+        Fix::safe_edit(Edit::range_replacement(replacement, stmt.range()))
+    } else {
+        // Ex) `import os, sys` (convert to `import os\nimport sys`)
+        let indentation = indentation_at_offset(stmt.start(), locator).unwrap_or_default();
+
+        // Generate newline-delimited imports.
+        let replacement = names
+            .iter()
+            .map(|alias| {
+                let Alias {
+                    range: _,
+                    name,
+                    asname,
+                } = alias;
+
+                if let Some(asname) = asname {
+                    format!("{indentation}import {name} as {asname}")
+                } else {
+                    format!("{indentation}import {name}")
+                }
+            })
+            .join(stylist.line_ending().as_str());
+
+        Fix::safe_edit(Edit::range_replacement(
+            replacement,
+            TextRange::new(locator.line_start(stmt.start()), stmt.end()),
+        ))
     }
 }

crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E401_E40.py.snap

@@ -6,8 +6,87 @@ E40.py:2:1: E401 Multiple imports on one line
 1 | #: E401
 2 | import os, sys
   | ^^^^^^^^^^^^^^ E401
-3 | #: Okay
-4 | import os
+3 | 
+4 | #: Okay
   |
+  = help: Split imports
+
+E40.py:65:1: E401 Multiple imports on one line
+   |
+64 | #: E401
+65 | import re as regex, string  # also with a comment!
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^ E401
+66 | import re as regex, string; x = 1
+   |
+   = help: Split imports
+
+E40.py:66:1: E401 Multiple imports on one line
+   |
+64 | #: E401
+65 | import re as regex, string  # also with a comment!
+66 | import re as regex, string; x = 1
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^ E401
+67 | 
+68 | x = 1; import re as regex, string
+   |
+   = help: Split imports
+
+E40.py:68:8: E401 Multiple imports on one line
+   |
+66 | import re as regex, string; x = 1
+67 | 
+68 | x = 1; import re as regex, string
+   |        ^^^^^^^^^^^^^^^^^^^^^^^^^^ E401
+   |
+   = help: Split imports
+
+E40.py:72:5: E401 Multiple imports on one line
+   |
+71 | def blah():
+72 |     import datetime as dt, copy
+   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^ E401
+73 | 
+74 |     def nested_and_tested():
+   |
+   = help: Split imports
+
+E40.py:75:9: E401 Multiple imports on one line
+   |
+74 |     def nested_and_tested():
+75 |         import builtins, textwrap as tw
+   |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E401
+76 | 
+77 |         x = 1; import re as regex, string
+   |
+   = help: Split imports
+
+E40.py:77:16: E401 Multiple imports on one line
+   |
+75 |         import builtins, textwrap as tw
+76 | 
+77 |         x = 1; import re as regex, string
+   |                ^^^^^^^^^^^^^^^^^^^^^^^^^^ E401
+78 |         import re as regex, string; x = 1
+   |
+   = help: Split imports
+
+E40.py:78:9: E401 Multiple imports on one line
+   |
+77 |         x = 1; import re as regex, string
+78 |         import re as regex, string; x = 1
+   |         ^^^^^^^^^^^^^^^^^^^^^^^^^^ E401
+79 | 
+80 |     if True: import re as regex, string
+   |
+   = help: Split imports
+
+E40.py:80:14: E401 Multiple imports on one line
+   |
+78 |         import re as regex, string; x = 1
+79 | 
+80 |     if True: import re as regex, string
+   |              ^^^^^^^^^^^^^^^^^^^^^^^^^^ E401
+   |
+   = help: Split imports
 
 

crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E402_E40.py.snap

@@ -1,32 +1,60 @@
 ---
 source: crates/ruff_linter/src/rules/pycodestyle/mod.rs
 ---
-E40.py:55:1: E402 Module level import not at top of file
+E40.py:56:1: E402 Module level import not at top of file
    |
-53 | VERSION = '1.2.3'
-54 | 
-55 | import foo
+54 | VERSION = '1.2.3'
+55 | 
+56 | import foo
    | ^^^^^^^^^^ E402
-56 | #: E402
-57 | import foo
+57 | #: E402
+58 | import foo
    |
 
-E40.py:57:1: E402 Module level import not at top of file
+E40.py:58:1: E402 Module level import not at top of file
    |
-55 | import foo
-56 | #: E402
-57 | import foo
+56 | import foo
+57 | #: E402
+58 | import foo
    | ^^^^^^^^^^ E402
-58 | 
-59 | a = 1
+59 | 
+60 | a = 1
    |
 
-E40.py:61:1: E402 Module level import not at top of file
+E40.py:62:1: E402 Module level import not at top of file
    |
-59 | a = 1
-60 | 
-61 | import bar
+60 | a = 1
+61 | 
+62 | import bar
    | ^^^^^^^^^^ E402
+63 | 
+64 | #: E401
+   |
+
+E40.py:65:1: E402 Module level import not at top of file
+   |
+64 | #: E401
+65 | import re as regex, string  # also with a comment!
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^ E402
+66 | import re as regex, string; x = 1
+   |
+
+E40.py:66:1: E402 Module level import not at top of file
+   |
+64 | #: E401
+65 | import re as regex, string  # also with a comment!
+66 | import re as regex, string; x = 1
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^ E402
+67 | 
+68 | x = 1; import re as regex, string
+   |
+
+E40.py:68:8: E402 Module level import not at top of file
+   |
+66 | import re as regex, string; x = 1
+67 | 
+68 | x = 1; import re as regex, string
+   |        ^^^^^^^^^^^^^^^^^^^^^^^^^^ E402
    |