Implement reimplemented_operator (FURB118)

crates/ruff_linter/resources/test/fixtures/refurb/FURB118.py

@@ -0,0 +1,56 @@
+# Errors.
+op_bitnot = lambda x: ~x
+op_not = lambda x: not x
+op_pos = lambda x: +x
+op_neg = lambda x: -x
+
+op_add = lambda x, y: x + y
+op_sub = lambda x, y: x - y
+op_mult = lambda x, y: x * y
+op_matmutl = lambda x, y: x @ y
+op_truediv = lambda x, y: x / y
+op_mod = lambda x, y: x % y
+op_pow = lambda x, y: x ** y
+op_lshift = lambda x, y: x << y
+op_rshift = lambda x, y: x >> y
+op_bitor = lambda x, y: x | y
+op_xor = lambda x, y: x ^ y
+op_bitand = lambda x, y: x & y
+op_floordiv = lambda x, y: x // y
+
+op_eq = lambda x, y: x == y
+op_ne = lambda x, y: x != y
+op_lt = lambda x, y: x < y
+op_lte = lambda x, y: x <= y
+op_gt = lambda x, y: x > y
+op_gte = lambda x, y: x >= y
+op_is = lambda x, y: x is y
+op_isnot = lambda x, y: x is not y
+op_in = lambda x, y: x in y
+
+def op_not2(x):
+    return not x
+
+def op_add2(x, y):
+    return x + y
+
+class Adder:
+    def add(x, y):
+        return x + y
+
+# Ok.
+op_add3 = lambda x, y = 1: x + y
+op_neg2 = lambda x, y: y - x
+op_notin = lambda x, y: y not in x
+op_and = lambda x, y: y and x
+op_or = lambda x, y: y or x
+
+def op_neg3(x, y):
+    return y - x
+
+def op_add4(x, y = 1):
+    return x + y
+
+def op_add5(x, y):
+    print("op_add5")
+    return x + y

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

@@ -2,7 +2,7 @@ use ruff_python_ast::Expr;
 
 use crate::checkers::ast::Checker;
 use crate::codes::Rule;
-use crate::rules::{flake8_pie, pylint};
+use crate::rules::{flake8_pie, pylint, refurb};
 
 /// Run lint rules over all deferred lambdas in the [`SemanticModel`].
 pub(crate) fn deferred_lambdas(checker: &mut Checker) {
@@ -21,6 +21,9 @@ pub(crate) fn deferred_lambdas(checker: &mut Checker) {
             if checker.enabled(Rule::ReimplementedContainerBuiltin) {
                 flake8_pie::rules::reimplemented_container_builtin(checker, lambda);
             }
+            if checker.enabled(Rule::ReimplementedOperator) {
+                refurb::rules::reimplemented_operator(checker, &lambda.into());
+            }
         }
     }
 }

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

@@ -368,6 +368,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                     .diagnostics
                     .extend(ruff::rules::unreachable::in_function(name, body));
             }
+            if checker.enabled(Rule::ReimplementedOperator) {
+                refurb::rules::reimplemented_operator(checker, &function_def.into());
+            }
         }
         Stmt::Return(_) => {
             if checker.enabled(Rule::ReturnOutsideFunction) {

crates/ruff_linter/src/codes.rs

@@ -952,6 +952,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Refurb, "105") => (RuleGroup::Preview, rules::refurb::rules::PrintEmptyString),
         #[allow(deprecated)]
         (Refurb, "113") => (RuleGroup::Nursery, rules::refurb::rules::RepeatedAppend),
+        (Refurb, "118") => (RuleGroup::Preview, rules::refurb::rules::ReimplementedOperator),
         #[allow(deprecated)]
         (Refurb, "131") => (RuleGroup::Nursery, rules::refurb::rules::DeleteFullSlice),
         #[allow(deprecated)]

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

@@ -16,6 +16,7 @@ mod tests {
 
     #[test_case(Rule::ReadWholeFile, Path::new("FURB101.py"))]
     #[test_case(Rule::RepeatedAppend, Path::new("FURB113.py"))]
+    #[test_case(Rule::ReimplementedOperator, Path::new("FURB118.py"))]
     #[test_case(Rule::DeleteFullSlice, Path::new("FURB131.py"))]
     #[test_case(Rule::CheckAndRemoveFromSet, Path::new("FURB132.py"))]
     #[test_case(Rule::IfExprMinMax, Path::new("FURB136.py"))]

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

@@ -8,6 +8,7 @@ pub(crate) use math_constant::*;
 pub(crate) use print_empty_string::*;
 pub(crate) use read_whole_file::*;
 pub(crate) use redundant_log_base::*;
+pub(crate) use reimplemented_operator::*;
 pub(crate) use reimplemented_starmap::*;
 pub(crate) use repeated_append::*;
 pub(crate) use single_item_membership_test::*;
@@ -25,6 +26,7 @@ mod math_constant;
 mod print_empty_string;
 mod read_whole_file;
 mod redundant_log_base;
+mod reimplemented_operator;
 mod reimplemented_starmap;
 mod repeated_append;
 mod single_item_membership_test;

crates/ruff_linter/src/rules/refurb/rules/reimplemented_operator.rs

@@ -0,0 +1,229 @@
+use anyhow::{bail, Result};
+use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
+use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::{self as ast, Expr, Stmt};
+use ruff_text_size::{Ranged, TextRange};
+
+use crate::checkers::ast::Checker;
+use crate::importer::ImportRequest;
+
+/// ## What it does
+/// Checks for lambda expressions and function definitions that can be replaced with a function
+/// from `operator` module.
+///
+/// ## Why is this bad?
+/// Using functions from `operator` module is more concise and readable.
+///
+/// ## Example
+/// ```python
+/// import functools
+///
+/// nums = [1, 2, 3]
+/// sum = functools.reduce(lambda x, y: x + y, nums)
+/// ```
+///
+/// Use instead:
+/// ```python
+/// import functools
+/// import operator
+///
+/// nums = [1, 2, 3]
+/// sum = functools.reduce(operator.add, nums)
+/// ```
+///
+/// ## References
+#[violation]
+pub struct ReimplementedOperator {
+    target: &'static str,
+    operator: &'static str,
+}
+
+impl Violation for ReimplementedOperator {
+    const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
+
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        let ReimplementedOperator { operator, target } = self;
+        format!("Use `operator.{operator}` instead of defining a {target}")
+    }
+
+    fn fix_title(&self) -> Option<String> {
+        let ReimplementedOperator { operator, .. } = self;
+        Some(format!("Replace with `operator.{operator}`"))
+    }
+}
+
+/// FURB118
+pub(crate) fn reimplemented_operator(checker: &mut Checker, target: &LambdaLike) {
+    let Some(params) = target.parameters() else {
+        return;
+    };
+    let Some(body) = target.body() else { return };
+    let Some(operator) = get_operator(body, params) else {
+        return;
+    };
+    let mut diagnostic = Diagnostic::new(
+        ReimplementedOperator {
+            operator,
+            target: target.kind(),
+        },
+        target.range(),
+    );
+    diagnostic.try_set_fix(|| target.try_fix(checker, operator));
+    checker.diagnostics.push(diagnostic);
+}
+
+/// Candidate for lambda expression or function definition consisting of a return statement.
+pub(crate) enum LambdaLike<'a> {
+    Lambda(&'a ast::ExprLambda),
+    Function(&'a ast::StmtFunctionDef),
+}
+
+impl<'a> From<&'a ast::ExprLambda> for LambdaLike<'a> {
+    fn from(lambda: &'a ast::ExprLambda) -> Self {
+        Self::Lambda(lambda)
+    }
+}
+
+impl<'a> From<&'a ast::StmtFunctionDef> for LambdaLike<'a> {
+    fn from(function: &'a ast::StmtFunctionDef) -> Self {
+        Self::Function(function)
+    }
+}
+
+impl Ranged for LambdaLike<'_> {
+    fn range(&self) -> TextRange {
+        match self {
+            Self::Lambda(expr) => expr.range(),
+            Self::Function(stmt) => stmt.range(),
+        }
+    }
+}
+
+impl LambdaLike<'_> {
+    fn parameters(&self) -> Option<&ast::Parameters> {
+        match self {
+            Self::Lambda(expr) => expr.parameters.as_deref(),
+            Self::Function(stmt) => Some(&stmt.parameters),
+        }
+    }
+
+    fn body(&self) -> Option<&Expr> {
+        match self {
+            Self::Lambda(expr) => Some(&expr.body),
+            Self::Function(stmt) => match stmt.body.as_slice() {
+                [Stmt::Return(ast::StmtReturn { value, .. })] => value.as_deref(),
+                _ => None,
+            },
+        }
+    }
+
+    fn try_fix(&self, checker: &Checker, operator: &'static str) -> Result<Fix> {
+        match self {
+            Self::Lambda(_) => {
+                let (edit, binding) = checker.importer().get_or_import_symbol(
+                    &ImportRequest::import("operator", operator),
+                    self.start(),
+                    checker.semantic(),
+                )?;
+                Ok(Fix::safe_edits(
+                    Edit::range_replacement(binding, self.range()),
+                    [edit],
+                ))
+            }
+            Self::Function(_) => bail!("No fix available"),
+        }
+    }
+
+    fn kind(&self) -> &'static str {
+        match self {
+            Self::Lambda(_) => "lambda",
+            Self::Function(_) => "function",
+        }
+    }
+}
+
+fn get_operator(expr: &Expr, params: &ast::Parameters) -> Option<&'static str> {
+    match expr {
+        Expr::UnaryOp(expr) => unary_op(expr, params),
+        Expr::BinOp(expr) => bin_op(expr, params),
+        Expr::Compare(expr) => cmp_op(expr, params),
+        _ => None,
+    }
+}
+
+fn unary_op(expr: &ast::ExprUnaryOp, params: &ast::Parameters) -> Option<&'static str> {
+    let [arg] = params.args.as_slice() else {
+        return None;
+    };
+    if !is_same(arg, &expr.operand) {
+        return None;
+    }
+    Some(match expr.op {
+        ast::UnaryOp::Invert => "invert",
+        ast::UnaryOp::Not => "not_",
+        ast::UnaryOp::UAdd => "pos",
+        ast::UnaryOp::USub => "neg",
+    })
+}
+
+fn bin_op(expr: &ast::ExprBinOp, params: &ast::Parameters) -> Option<&'static str> {
+    let [arg1, arg2] = params.args.as_slice() else {
+        return None;
+    };
+    if !is_same(arg1, &expr.left) || !is_same(arg2, &expr.right) {
+        return None;
+    }
+    Some(match expr.op {
+        ast::Operator::Add => "add",
+        ast::Operator::Sub => "sub",
+        ast::Operator::Mult => "mul",
+        ast::Operator::MatMult => "matmul",
+        ast::Operator::Div => "truediv",
+        ast::Operator::Mod => "mod",
+        ast::Operator::Pow => "pow",
+        ast::Operator::LShift => "lshift",
+        ast::Operator::RShift => "rshift",
+        ast::Operator::BitOr => "or_",
+        ast::Operator::BitXor => "xor",
+        ast::Operator::BitAnd => "and_",
+        ast::Operator::FloorDiv => "floordiv",
+    })
+}
+
+fn cmp_op(expr: &ast::ExprCompare, params: &ast::Parameters) -> Option<&'static str> {
+    let [arg1, arg2] = params.args.as_slice() else {
+        return None;
+    };
+    let [op] = expr.ops.as_slice() else {
+        return None;
+    };
+    let [right] = expr.comparators.as_slice() else {
+        return None;
+    };
+    if !is_same(arg1, &expr.left) || !is_same(arg2, right) {
+        return None;
+    }
+    match op {
+        ast::CmpOp::Eq => Some("eq"),
+        ast::CmpOp::NotEq => Some("ne"),
+        ast::CmpOp::Lt => Some("lt"),
+        ast::CmpOp::LtE => Some("le"),
+        ast::CmpOp::Gt => Some("gt"),
+        ast::CmpOp::GtE => Some("ge"),
+        ast::CmpOp::Is => Some("is_"),
+        ast::CmpOp::IsNot => Some("is_not"),
+        ast::CmpOp::In => Some("contains"),
+        ast::CmpOp::NotIn => None,
+    }
+}
+
+fn is_same(arg: &ast::ParameterWithDefault, expr: &Expr) -> bool {
+    if arg.default.is_some() {
+        false
+    } else if let Expr::Name(name) = expr {
+        name.id == arg.parameter.name.as_str()
+    } else {
+        false
+    }
+}