[pylint] Preserve original value format (PLR6104)

crates/ruff_linter/resources/test/fixtures/pylint/non_augmented_assignment.py

@@ -50,6 +50,50 @@ def t(self):
 obj = T()
 obj.a = obj.a + 1
 
+
+a = a+-1
+
+# Regression tests for https://github.com/astral-sh/ruff/issues/11672
+test = 0x5
+test = test + 0xBA
+
+test2 = b""
+test2 = test2 + b"\000"
+
+test3 = ""
+test3 = test3 + (   a := R""
+                         f"oo"   )
+
+test4 = []
+test4 = test4 + ( e
+                  for e in
+                  range(10)
+                  )
+
+test5 = test5 + (
+    4
+    *
+    10
+)
+
+test6 = test6 + \
+        (
+            4
+            *
+            10
+        )
+
+test7 = \
+        100 \
+    + test7
+
+test8 = \
+    886 \
+    + \
+ \
+    test8
+
+
 # OK
 a_list[0] = a_list[:] * 3
 index = a_number = a_number + 1

crates/ruff_linter/src/rules/pylint/rules/non_augmented_assignment.rs

@@ -1,13 +1,14 @@
-use ast::{Expr, StmtAugAssign};
+use ast::Expr;
 use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
 use ruff_python_ast as ast;
 use ruff_python_ast::comparable::ComparableExpr;
 use ruff_python_ast::Operator;
-use ruff_python_codegen::Generator;
+use ruff_python_trivia::Cursor;
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::checkers::ast::Checker;
+use crate::Locator;
 
 /// ## What it does
 /// Checks for assignments that can be replaced with augmented assignment
@@ -97,16 +98,20 @@ pub(crate) fn non_augmented_assignment(checker: &mut Checker, assign: &ast::Stmt
         return;
     };
 
+    let locator = checker.locator();
     let operator = AugmentedOperator::from(value.op);
 
     // Match, e.g., `x = x + 1`.
     if ComparableExpr::from(target) == ComparableExpr::from(&value.left) {
         let mut diagnostic = Diagnostic::new(NonAugmentedAssignment { operator }, assign.range());
+        let half_expr = locator.slice(TextRange::new(value.left.end(), value.end()));
+        let right_operand_expr = trim_left_operator(half_expr);
+
         diagnostic.set_fix(Fix::unsafe_edit(augmented_assignment(
-            checker.generator(),
+            checker.locator(),
             target,
-            value.op,
-            &value.right,
+            operator,
+            right_operand_expr,
             assign.range(),
         )));
         checker.diagnostics.push(diagnostic);
@@ -120,36 +125,98 @@ pub(crate) fn non_augmented_assignment(checker: &mut Checker, assign: &ast::Stmt
         && ComparableExpr::from(target) == ComparableExpr::from(&value.right)
     {
         let mut diagnostic = Diagnostic::new(NonAugmentedAssignment { operator }, assign.range());
+        let half_expr = locator.slice(TextRange::new(value.left.start(), value.right.start()));
+        let right_operand_expr = trim_right_operator(half_expr);
+
         diagnostic.set_fix(Fix::unsafe_edit(augmented_assignment(
-            checker.generator(),
+            checker.locator(),
             target,
-            value.op,
-            &value.left,
+            operator,
+            right_operand_expr,
             assign.range(),
         )));
         checker.diagnostics.push(diagnostic);
     }
 }
 
+const OPERATORS: [&str; 15] = [
+    "+", "&", "|", "^", "//", "/", "<<", "<", "@", "%", "**", "*", ">>", ">", "-",
+];
+
+macro_rules! trim_operator {
+    ($half_expr:ident, $eat:ident, $bump:ident, $check:ident) => {{
+        let mut cursor = Cursor::new($half_expr);
+
+        cursor.$eat(is_whitespace_or_line_continuation);
+
+        let op = OPERATORS
+            .iter()
+            .find(|&op| cursor.as_str().$check(op))
+            .unwrap();
+
+        cursor.$bump();
+
+        if op.len() == 2 {
+            cursor.$bump();
+        }
+
+        cursor.$eat(is_whitespace_or_line_continuation);
+
+        cursor.as_str()
+    }};
+}
+
+/// Input:
+/// ```python
+///   \
+///    + (a -\
+/// b)
+/// ```
+///
+/// Output:
+/// ```python
+/// (a -\
+/// b)
+/// ```
+fn trim_left_operator(half_expr: &str) -> &str {
+    trim_operator!(half_expr, eat_while, bump, starts_with)
+}
+
+/// Input:
+/// ```python
+/// (a -\
+/// b) \
+///     +  \
+///
+/// ```
+///
+/// Output:
+/// ```python
+/// (a -\
+/// b)
+/// ```
+fn trim_right_operator(half_expr: &str) -> &str {
+    trim_operator!(half_expr, eat_back_while, bump_back, ends_with)
+}
+
+fn is_whitespace_or_line_continuation(c: char) -> bool {
+    c == '\\' || c.is_whitespace()
+}
+
 /// Generate a fix to convert an assignment statement to an augmented assignment.
 ///
 /// For example, given `x = x + 1`, the fix would be `x += 1`.
 fn augmented_assignment(
-    generator: Generator,
+    locator: &Locator,
     target: &Expr,
-    operator: Operator,
-    right_operand: &Expr,
+    operator: AugmentedOperator,
+    right_operand_expr: &str,
     range: TextRange,
 ) -> Edit {
-    Edit::range_replacement(
-        generator.stmt(&ast::Stmt::AugAssign(StmtAugAssign {
-            range: TextRange::default(),
-            target: Box::new(target.clone()),
-            op: operator,
-            value: Box::new(right_operand.clone()),
-        })),
-        range,
-    )
+    let target_expr = locator.slice(target);
+    let new_content = format!("{target_expr} {operator} {right_operand_expr}");
+
+    Edit::range_replacement(new_content, range)
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]