[pylint] Implement rule to prefer augmented assignment (PLR6104)

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

@@ -0,0 +1,55 @@
+# Errors
+some_string = "some string"
+index, a_number, to_multiply, to_divide, to_cube, timeDiffSeconds, flags = (
+    0,
+    1,
+    2,
+    3,
+    4,
+    5,
+    0x3,
+)
+a_list = [1, 2]
+some_set = {"elem"}
+mat1, mat2 = None, None
+
+some_string = some_string + "a very long end of string"
+index = index - 1
+a_list = a_list + ["to concat"]
+some_set = some_set | {"to concat"}
+to_multiply = to_multiply * 5
+to_divide = to_divide / 5
+to_divide = to_divide // 5
+to_cube = to_cube**3
+to_cube = 3**to_cube
+to_cube = to_cube**to_cube
+timeDiffSeconds = timeDiffSeconds % 60
+flags = flags & 0x1
+flags = flags | 0x1
+flags = flags ^ 0x1
+flags = flags << 1
+flags = flags >> 1
+mat1 = mat1 @ mat2
+a_list[1] = a_list[1] + 1
+
+a_list[0:2] = a_list[0:2] * 3
+a_list[:2] = a_list[:2] * 3
+a_list[1:] = a_list[1:] * 3
+a_list[:] = a_list[:] * 3
+
+index = index * (index + 10)
+
+
+class T:
+    def t(self):
+        self.a = self.a + 1
+
+
+obj = T()
+obj.a = obj.a + 1
+
+# OK
+a_list[0] = a_list[:] * 3
+index = a_number = a_number + 1
+a_number = index = a_number + 1
+index = index * index + 10

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -1478,6 +1478,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.settings.rules.enabled(Rule::TypeBivariance) {
                 pylint::rules::type_bivariance(checker, value);
             }
+            if checker.enabled(Rule::NonAugmentedAssignment) {
+                pylint::rules::non_augmented_assignment(checker, assign);
+            }
             if checker.settings.rules.enabled(Rule::UnsortedDunderAll) {
                 ruff::rules::sort_dunder_all_assign(checker, assign);
             }

crates/ruff_linter/src/codes.rs

@@ -293,6 +293,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Pylint, "R2004") => (RuleGroup::Stable, rules::pylint::rules::MagicValueComparison),
         (Pylint, "R2044") => (RuleGroup::Preview, rules::pylint::rules::EmptyComment),
         (Pylint, "R5501") => (RuleGroup::Stable, rules::pylint::rules::CollapsibleElseIf),
+        (Pylint, "R6104") => (RuleGroup::Preview, rules::pylint::rules::NonAugmentedAssignment),
         (Pylint, "R6201") => (RuleGroup::Preview, rules::pylint::rules::LiteralMembership),
         #[allow(deprecated)]
         (Pylint, "R6301") => (RuleGroup::Nursery, rules::pylint::rules::NoSelfUse),

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

@@ -186,6 +186,7 @@ mod tests {
         Rule::UnnecessaryDictIndexLookup,
         Path::new("unnecessary_dict_index_lookup.py")
     )]
+    #[test_case(Rule::NonAugmentedAssignment, Path::new("non_augmented_assignment.py"))]
     #[test_case(
         Rule::UselessExceptionStatement,
         Path::new("useless_exception_statement.py")

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

@@ -47,6 +47,7 @@ pub(crate) use no_method_decorator::*;
 pub(crate) use no_self_use::*;
 pub(crate) use non_ascii_module_import::*;
 pub(crate) use non_ascii_name::*;
+pub(crate) use non_augmented_assignment::*;
 pub(crate) use non_slot_assignment::*;
 pub(crate) use nonlocal_and_global::*;
 pub(crate) use nonlocal_without_binding::*;
@@ -143,6 +144,7 @@ mod no_method_decorator;
 mod no_self_use;
 mod non_ascii_module_import;
 mod non_ascii_name;
+mod non_augmented_assignment;
 mod non_slot_assignment;
 mod nonlocal_and_global;
 mod nonlocal_without_binding;

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

@@ -0,0 +1,202 @@
+use ast::{Expr, StmtAugAssign};
+use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
+use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast as ast;
+use ruff_python_ast::comparable::ComparableExpr;
+use ruff_python_ast::Operator;
+use ruff_python_codegen::Generator;
+use ruff_text_size::{Ranged, TextRange};
+
+use crate::checkers::ast::Checker;
+
+/// ## What it does
+/// Checks for assignments that can be replaced with augmented assignment
+/// statements.
+///
+/// ## Why is this bad?
+/// If an assignment statement consists of a binary operation in which one
+/// operand is the same as the assignment target, it can be rewritten as an
+/// augmented assignment. For example, `x = x + 1` can be rewritten as
+/// `x += 1`.
+///
+/// When performing such an operation, augmented assignments are more concise
+/// and idiomatic.
+///
+/// ## Example
+/// ```python
+/// x = x + 1
+/// ```
+///
+/// Use instead:
+/// ```python
+/// x += 1
+/// ```
+///
+/// ## Fix safety
+/// This rule's fix is marked as unsafe, as augmented assignments have
+/// different semantics when the target is a mutable data type, like a list or
+/// dictionary.
+///
+/// For example, consider the following:
+///
+/// ```python
+/// foo = [1]
+/// bar = foo
+/// foo = foo + [2]
+/// assert (foo, bar) == ([1, 2], [1])
+/// ```
+///
+/// If the assignment is replaced with an augmented assignment, the update
+/// operation will apply to both `foo` and `bar`, as they refer to the same
+/// object:
+///
+/// ```python
+/// foo = [1]
+/// bar = foo
+/// foo += [2]
+/// assert (foo, bar) == ([1, 2], [1, 2])
+/// ```
+#[violation]
+pub struct NonAugmentedAssignment {
+    operator: AugmentedOperator,
+}
+
+impl AlwaysFixableViolation for NonAugmentedAssignment {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        let NonAugmentedAssignment { operator } = self;
+        format!("Use `{operator}` to perform an augmented assignment directly")
+    }
+
+    fn fix_title(&self) -> String {
+        "Replace with augmented assignment".to_string()
+    }
+}
+
+/// PLR6104
+pub(crate) fn non_augmented_assignment(checker: &mut Checker, assign: &ast::StmtAssign) {
+    // Ignore multiple assignment targets.
+    let [target] = assign.targets.as_slice() else {
+        return;
+    };
+
+    // Match, e.g., `x = x + 1`.
+    let Expr::BinOp(value) = &*assign.value else {
+        return;
+    };
+
+    // Match, e.g., `x = x + 1`.
+    if ComparableExpr::from(target) == ComparableExpr::from(&value.left) {
+        let mut diagnostic = Diagnostic::new(
+            NonAugmentedAssignment {
+                operator: AugmentedOperator::from(value.op),
+            },
+            assign.range(),
+        );
+        diagnostic.set_fix(Fix::unsafe_edit(augmented_assignment(
+            checker.generator(),
+            target,
+            value.op,
+            &value.right,
+            assign.range(),
+        )));
+        checker.diagnostics.push(diagnostic);
+        return;
+    }
+
+    // Match, e.g., `x = 1 + x`.
+    if ComparableExpr::from(target) == ComparableExpr::from(&value.right) {
+        let mut diagnostic = Diagnostic::new(
+            NonAugmentedAssignment {
+                operator: AugmentedOperator::from(value.op),
+            },
+            assign.range(),
+        );
+        diagnostic.set_fix(Fix::unsafe_edit(augmented_assignment(
+            checker.generator(),
+            target,
+            value.op,
+            &value.left,
+            assign.range(),
+        )));
+        checker.diagnostics.push(diagnostic);
+    }
+}
+
+/// 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,
+    target: &Expr,
+    operator: Operator,
+    right_operand: &Expr,
+    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,
+    )
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum AugmentedOperator {
+    Add,
+    BitAnd,
+    BitOr,
+    BitXor,
+    Div,
+    FloorDiv,
+    LShift,
+    MatMult,
+    Mod,
+    Mult,
+    Pow,
+    RShift,
+    Sub,
+}
+
+impl From<Operator> for AugmentedOperator {
+    fn from(value: Operator) -> Self {
+        match value {
+            Operator::Add => Self::Add,
+            Operator::BitAnd => Self::BitAnd,
+            Operator::BitOr => Self::BitOr,
+            Operator::BitXor => Self::BitXor,
+            Operator::Div => Self::Div,
+            Operator::FloorDiv => Self::FloorDiv,
+            Operator::LShift => Self::LShift,
+            Operator::MatMult => Self::MatMult,
+            Operator::Mod => Self::Mod,
+            Operator::Mult => Self::Mult,
+            Operator::Pow => Self::Pow,
+            Operator::RShift => Self::RShift,
+            Operator::Sub => Self::Sub,
+        }
+    }
+}
+
+impl std::fmt::Display for AugmentedOperator {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Add => f.write_str("+="),
+            Self::BitAnd => f.write_str("&="),
+            Self::BitOr => f.write_str("|="),
+            Self::BitXor => f.write_str("^="),
+            Self::Div => f.write_str("/="),
+            Self::FloorDiv => f.write_str("//="),
+            Self::LShift => f.write_str("<<="),
+            Self::MatMult => f.write_str("@="),
+            Self::Mod => f.write_str("%="),
+            Self::Mult => f.write_str("*="),
+            Self::Pow => f.write_str("**="),
+            Self::RShift => f.write_str(">>="),
+            Self::Sub => f.write_str("-="),
+        }
+    }
+}