Autofix for none-not-at-end-of-union (RUF036)

crates/ruff_linter/resources/test/fixtures/ruff/RUF036.py

@@ -25,6 +25,14 @@ def func6(arg: U[None, None, int]):
     ...
 
 
+def func7() -> U[
+    None,
+    # comment
+    int
+]: 
+    ...
+
+
 # Ok
 def good_func1(arg: int | None): 
     ...

crates/ruff_linter/src/rules/ruff/rules/none_not_at_end_of_union.rs

@@ -1,6 +1,7 @@
-use ruff_diagnostics::{Diagnostic, Violation};
+use itertools::{Itertools, Position};
+use ruff_diagnostics::{Applicability, Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
-use ruff_python_ast::Expr;
+use ruff_python_ast::{self as ast, Expr};
 use ruff_python_semantic::analyze::typing::traverse_union;
 use ruff_text_size::Ranged;
 use smallvec::SmallVec;
@@ -33,21 +34,29 @@ use crate::checkers::ast::Checker;
 pub(crate) struct NoneNotAtEndOfUnion;
 
 impl Violation for NoneNotAtEndOfUnion {
+    const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
+
     #[derive_message_formats]
     fn message(&self) -> String {
         "`None` not at the end of the type annotation.".to_string()
     }
+
+    fn fix_title(&self) -> Option<String> {
+        Some("Move `None` to the end of the type annotation".to_string())
+    }
 }
 
 /// RUF036
 pub(crate) fn none_not_at_end_of_union<'a>(checker: &mut Checker, union: &'a Expr) {
     let semantic = checker.semantic();
     let mut none_exprs: SmallVec<[&Expr; 1]> = SmallVec::new();
-
+    let mut non_none_exprs: SmallVec<[&Expr; 1]> = SmallVec::new();
     let mut last_expr: Option<&Expr> = None;
     let mut find_none = |expr: &'a Expr, _parent: &Expr| {
         if matches!(expr, Expr::NoneLiteral(_)) {
             none_exprs.push(expr);
+        } else {
+            non_none_exprs.push(expr);
         }
         last_expr = Some(expr);
     };
@@ -59,7 +68,7 @@ pub(crate) fn none_not_at_end_of_union<'a>(checker: &mut Checker, union: &'a Exp
         return;
     };
 
-    // The must be at least one `None` expression.
+    // There must be at least one `None` expression.
     let Some(last_none) = none_exprs.last() else {
         return;
     };
@@ -69,9 +78,37 @@ pub(crate) fn none_not_at_end_of_union<'a>(checker: &mut Checker, union: &'a Exp
         return;
     }
 
-    for none_expr in none_exprs {
-        checker
-            .diagnostics
-            .push(Diagnostic::new(NoneNotAtEndOfUnion, none_expr.range()));
+    for (pos, none_expr) in none_exprs.iter().with_position() {
+        let mut diagnostic = Diagnostic::new(NoneNotAtEndOfUnion, none_expr.range());
+        if matches!(pos, Position::Last | Position::Only) {
+            let mut elements = non_none_exprs
+                .iter()
+                .map(|expr| checker.locator().slice(expr.range()).to_string())
+                .chain(std::iter::once("None".to_string()));
+            let fix = if let Expr::Subscript(ast::ExprSubscript { slice, .. }) = union {
+                let applicability = if checker.comment_ranges().intersects(slice.range()) {
+                    Applicability::Unsafe
+                } else {
+                    Applicability::Safe
+                };
+                Fix::applicable_edit(
+                    Edit::range_replacement(elements.join(", "), slice.range()),
+                    applicability,
+                )
+            } else {
+                let applicability = if checker.comment_ranges().intersects(union.range()) {
+                    Applicability::Unsafe
+                } else {
+                    Applicability::Safe
+                };
+                Fix::applicable_edit(
+                    Edit::range_replacement(elements.join(" | "), union.range()),
+                    applicability,
+                )
+            };
+            diagnostic.set_fix(fix);
+        }
+
+        checker.diagnostics.push(diagnostic);
     }
 }

crates/ruff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF036_RUF036.py.snap

@@ -2,58 +2,148 @@
 source: crates/ruff_linter/src/rules/ruff/mod.rs
 snapshot_kind: text
 ---
-RUF036.py:4:16: RUF036 `None` not at the end of the type annotation.
+RUF036.py:4:16: RUF036 [*] `None` not at the end of the type annotation.
   |
 4 | def func1(arg: None | int): 
   |                ^^^^ RUF036
 5 |     ...
   |
+  = help: Move `None` to the end of the type annotation
 
-RUF036.py:8:16: RUF036 `None` not at the end of the type annotation.
+ℹ Safe fix
+1 1 | from typing import Union as U
+2 2 | 
+3 3 | 
+4   |-def func1(arg: None | int): 
+  4 |+def func1(arg: int | None): 
+5 5 |     ...
+6 6 | 
+7 7 | 
+
+RUF036.py:8:16: RUF036 [*] `None` not at the end of the type annotation.
   |
 8 | def func2() -> None | int: 
   |                ^^^^ RUF036
 9 |     ...
   |
+  = help: Move `None` to the end of the type annotation
+
+ℹ Safe fix
+5 5 |     ...
+6 6 | 
+7 7 | 
+8   |-def func2() -> None | int: 
+  8 |+def func2() -> int | None: 
+9 9 |     ...
+10 10 | 
+11 11 | 
 
 RUF036.py:12:16: RUF036 `None` not at the end of the type annotation.
    |
 12 | def func3(arg: None | None | int): 
    |                ^^^^ RUF036
 13 |     ...
    |
+   = help: Move `None` to the end of the type annotation
 
-RUF036.py:12:23: RUF036 `None` not at the end of the type annotation.
+RUF036.py:12:23: RUF036 [*] `None` not at the end of the type annotation.
    |
 12 | def func3(arg: None | None | int): 
    |                       ^^^^ RUF036
 13 |     ...
    |
+   = help: Move `None` to the end of the type annotation
+
+ℹ Safe fix
+9  9  |     ...
+10 10 | 
+11 11 | 
+12    |-def func3(arg: None | None | int): 
+   12 |+def func3(arg: int | None): 
+13 13 |     ...
+14 14 | 
+15 15 | 
 
-RUF036.py:16:18: RUF036 `None` not at the end of the type annotation.
+RUF036.py:16:18: RUF036 [*] `None` not at the end of the type annotation.
    |
 16 | def func4(arg: U[None, int]): 
    |                  ^^^^ RUF036
 17 |     ...
    |
+   = help: Move `None` to the end of the type annotation
 
-RUF036.py:20:18: RUF036 `None` not at the end of the type annotation.
+ℹ Safe fix
+13 13 |     ...
+14 14 | 
+15 15 | 
+16    |-def func4(arg: U[None, int]): 
+   16 |+def func4(arg: U[int, None]): 
+17 17 |     ...
+18 18 | 
+19 19 | 
+
+RUF036.py:20:18: RUF036 [*] `None` not at the end of the type annotation.
    |
 20 | def func5() -> U[None, int]: 
    |                  ^^^^ RUF036
 21 |     ...
    |
+   = help: Move `None` to the end of the type annotation
+
+ℹ Safe fix
+17 17 |     ...
+18 18 | 
+19 19 | 
+20    |-def func5() -> U[None, int]: 
+   20 |+def func5() -> U[int, None]: 
+21 21 |     ...
+22 22 | 
+23 23 | 
 
 RUF036.py:24:18: RUF036 `None` not at the end of the type annotation.
    |
 24 | def func6(arg: U[None, None, int]): 
    |                  ^^^^ RUF036
 25 |     ...
    |
+   = help: Move `None` to the end of the type annotation
 
-RUF036.py:24:24: RUF036 `None` not at the end of the type annotation.
+RUF036.py:24:24: RUF036 [*] `None` not at the end of the type annotation.
    |
 24 | def func6(arg: U[None, None, int]): 
    |                        ^^^^ RUF036
 25 |     ...
    |
+   = help: Move `None` to the end of the type annotation
+
+ℹ Safe fix
+21 21 |     ...
+22 22 | 
+23 23 | 
+24    |-def func6(arg: U[None, None, int]): 
+   24 |+def func6(arg: U[int, None]): 
+25 25 |     ...
+26 26 | 
+27 27 | 
+
+RUF036.py:29:5: RUF036 [*] `None` not at the end of the type annotation.
+   |
+28 | def func7() -> U[
+29 |     None,
+   |     ^^^^ RUF036
+30 |     # comment
+31 |     int
+   |
+   = help: Move `None` to the end of the type annotation
+
+ℹ Unsafe fix
+26 26 | 
+27 27 | 
+28 28 | def func7() -> U[
+29    |-    None,
+30    |-    # comment
+31    |-    int
+   29 |+    int, None
+32 30 | ]: 
+33 31 |     ...
+34 32 |