[pylint] - add fix for unary expressions in PLC2801

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

@@ -1,6 +1,6 @@
 from typing import Any
 
-
+a = 2
 print((3.0).__add__(4.0))  # PLC2801
 print((3.0).__sub__(4.0))  # PLC2801
 print((3.0).__mul__(4.0))  # PLC2801
@@ -17,6 +17,67 @@
 print((3.0).__repr__())  # PLC2801
 print([1, 2, 3].__len__())  # PLC2801
 print((1).__neg__())  # PLC2801
+print(-a.__sub__(1))  # PLC2801
+print(-(a).__sub__(1))  # PLC2801
+print(-(-a.__sub__(1)))  # PLC2801
+print((5 - a).__sub__(1))  # PLC2801
+print(-(5 - a).__sub__(1))  # PLC2801
+print(-(-5 - a).__sub__(1))  # PLC2801
+print(+-+-+-a.__sub__(1))  # PLC2801
+print(a.__rsub__(2 - 1))  # PLC2801
+print(a.__sub__(((((1))))))  # PLC2801
+print(a.__sub__(((((2 - 1))))))  # PLC2801
+print(a.__sub__(
+    3
+    +
+    4
+))
+print(a.__rsub__(
+    3
+    +
+    4
+))
+print(2 * a.__add__(3))  # PLC2801
+x = 2 * a.__add__(3)  # PLC2801
+x = 2 * -a.__add__(3)  # PLC2801
+x = a.__add__(3)  # PLC2801
+x = -a.__add__(3)  # PLC2801
+x = (-a).__add__(3)  # PLC2801
+x = -(-a).__add__(3)  # PLC2801
+
+# Calls
+print(a.__call__())  # PLC2801 (no fix, intentional)
+
+# Lambda expressions
+blah = lambda: a.__add__(1)  # PLC2801
+
+# If expressions
+print(a.__add__(1) if a > 0 else a.__sub__(1))  # PLC2801
+
+# Dict/Set/List/Tuple
+print({"a": a.__add__(1)})  # PLC2801
+print({a.__add__(1)})  # PLC2801
+print([a.__add__(1)])  # PLC2801
+print((a.__add__(1),))  # PLC2801
+
+# Comprehension variants
+print({i: i.__add__(1) for i in range(5)})  # PLC2801
+print({i.__add__(1) for i in range(5)})  # PLC2801
+print([i.__add__(1) for i in range(5)])  # PLC2801
+print((i.__add__(1) for i in range(5)))  # PLC2801
+
+# Generators
+gen = (i.__add__(1) for i in range(5))  # PLC2801
+print(next(gen))
+
+# Subscripts
+print({"a": a.__add__(1)}["a"])  # PLC2801
+
+# Starred
+print(*[a.__add__(1)])  # PLC2801
+
+# Slices
+print([a.__add__(1), a.__sub__(1)][0:1])  # PLC2801
 
 
 class Thing:

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

@@ -124,21 +124,30 @@ pub(crate) fn unnecessary_dunder_call(checker: &mut Checker, call: &ast::ExprCal
                 title = Some(message.to_string());
             }
             ([arg], DunderReplacement::Operator(replacement, message)) => {
-                fixed = Some(format!(
-                    "{} {} {}",
-                    checker.locator().slice(value.as_ref()),
-                    replacement,
-                    checker.locator().slice(arg),
-                ));
+                let value_slice = checker.locator().slice(value.as_ref());
+                let arg_slice = checker.locator().slice(arg);
+
+                if can_be_represented_without_parentheses(arg) {
+                    // if it's something that can reasonably be removed from parentheses,
+                    // we'll do that.
+                    fixed = Some(format!("{value_slice} {replacement} {arg_slice}"));
+                } else {
+                    fixed = Some(format!("{value_slice} {replacement} ({arg_slice})"));
+                }
+
                 title = Some(message.to_string());
             }
             ([arg], DunderReplacement::ROperator(replacement, message)) => {
-                fixed = Some(format!(
-                    "{} {} {}",
-                    checker.locator().slice(arg),
-                    replacement,
-                    checker.locator().slice(value.as_ref()),
-                ));
+                let value_slice = checker.locator().slice(value.as_ref());
+                let arg_slice = checker.locator().slice(arg);
+
+                if arg.is_attribute_expr() || arg.is_name_expr() || arg.is_literal_expr() {
+                    // if it's something that can reasonably be removed from parentheses,
+                    // we'll do that.
+                    fixed = Some(format!("{arg_slice} {replacement} {value_slice}"));
+                } else {
+                    fixed = Some(format!("({arg_slice}) {replacement} {value_slice}"));
+                }
                 title = Some(message.to_string());
             }
             (_, DunderReplacement::MessageOnly(message)) => {
@@ -156,8 +165,25 @@ pub(crate) fn unnecessary_dunder_call(checker: &mut Checker, call: &ast::ExprCal
         call.range(),
     );
 
-    if let Some(fixed) = fixed {
-        diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(fixed, call.range())));
+    if let Some(mut fixed) = fixed {
+        // by the looks of it, we don't need to wrap the expression in parens if
+        // it's the only argument to a call expression.
+        // being in any other kind of expression though, we *will* add parens.
+        // e.g. `print(a.__add__(3))` -> `print(a + 3)` instead of `print((a + 3))`
+        // a multiplication expression example: `x = 2 * a.__add__(3)` -> `x = 2 * (a + 3)`
+        let wrap_in_paren = checker
+            .semantic()
+            .current_expression_parent()
+            .is_some_and(|parent| !can_be_represented_without_parentheses(parent));
+
+        if wrap_in_paren {
+            fixed = format!("({fixed})");
+        }
+
+        diagnostic.set_fix(Fix::unsafe_edit(Edit::range_replacement(
+            fixed,
+            call.range(),
+        )));
     };
 
     checker.diagnostics.push(diagnostic);
@@ -358,3 +384,24 @@ fn in_dunder_method_definition(semantic: &SemanticModel) -> bool {
         func_def.name.starts_with("__") && func_def.name.ends_with("__")
     })
 }
+
+/// Returns `true` if the [`Expr`] can be represented without parentheses.
+fn can_be_represented_without_parentheses(expr: &Expr) -> bool {
+    expr.is_attribute_expr()
+        || expr.is_name_expr()
+        || expr.is_literal_expr()
+        || expr.is_call_expr()
+        || expr.is_lambda_expr()
+        || expr.is_if_exp_expr()
+        || expr.is_generator_exp_expr()
+        || expr.is_subscript_expr()
+        || expr.is_starred_expr()
+        || expr.is_slice_expr()
+        || expr.is_dict_expr()
+        || expr.is_dict_comp_expr()
+        || expr.is_list_expr()
+        || expr.is_list_comp_expr()
+        || expr.is_tuple_expr()
+        || expr.is_set_comp_expr()
+        || expr.is_set_expr()
+}