[refurb] Support itemgetter in reimplemented-operator (FURB118)

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

@@ -10,7 +10,7 @@
 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_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
@@ -27,6 +27,10 @@
 op_is = lambda x, y: x is y
 op_isnot = lambda x, y: x is not y
 op_in = lambda x, y: y in x
+op_itemgetter = lambda x: x[0]
+op_itemgetter = lambda x: (x[0], x[1], x[2])
+op_itemgetter = lambda x: (x[1:], x[2])
+op_itemgetter = lambda x: x[:]
 
 
 def op_not2(x):
@@ -41,21 +45,30 @@ class Adder:
     def add(x, y):
         return x + y
 
+
 # OK.
-op_add3 = lambda x, y = 1: x + y
+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
 op_in = lambda x, y: x in y
+op_itemgetter = lambda x: (1, x[1], x[2])
+op_itemgetter = lambda x: (x.y, x[1], x[2])
+op_itemgetter = lambda x, y: (x[0], y[0])
+op_itemgetter = lambda x, y: (x[0], y[0])
+op_itemgetter = lambda x: ()
+op_itemgetter = lambda x: (*x[0], x[1])
 
 
 def op_neg3(x, y):
     return y - x
 
-def op_add4(x, y = 1):
+
+def op_add4(x, y=1):
     return x + y
 
+
 def op_add5(x, y):
     print("op_add5")
     return x + y

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

@@ -1,9 +1,13 @@
+use std::borrow::Cow;
+use std::fmt::{Debug, Display, Formatter};
+
 use anyhow::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_python_ast::{self as ast, Expr, ExprSlice, ExprSubscript, ExprTuple, Parameters, Stmt};
 use ruff_python_semantic::SemanticModel;
+use ruff_source_file::Locator;
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::checkers::ast::Checker;
@@ -39,7 +43,7 @@ use crate::importer::{ImportRequest, Importer};
 /// ## References
 #[violation]
 pub struct ReimplementedOperator {
-    operator: &'static str,
+    operator: Operator,
     target: FunctionLikeKind,
 }
 
@@ -71,18 +75,18 @@ pub(crate) fn reimplemented_operator(checker: &mut Checker, target: &FunctionLik
         return;
     };
     let Some(body) = target.body() else { return };
-    let Some(operator) = get_operator(body, params) else {
+    let Some(operator) = get_operator(body, params, checker.locator()) else {
         return;
     };
+    let fix = target.try_fix(&operator, checker.importer(), checker.semantic());
     let mut diagnostic = Diagnostic::new(
         ReimplementedOperator {
             operator,
             target: target.kind(),
         },
         target.range(),
     );
-    diagnostic
-        .try_set_optional_fix(|| target.try_fix(operator, checker.importer(), checker.semantic()));
+    diagnostic.try_set_optional_fix(|| fix);
     checker.diagnostics.push(diagnostic);
 }
 
@@ -115,8 +119,8 @@ impl Ranged for FunctionLike<'_> {
 }
 
 impl FunctionLike<'_> {
-    /// Return the [`ast::Parameters`] of the function-like node.
-    fn parameters(&self) -> Option<&ast::Parameters> {
+    /// Return the [`Parameters`] of the function-like node.
+    fn parameters(&self) -> Option<&Parameters> {
         match self {
             Self::Lambda(expr) => expr.parameters.as_deref(),
             Self::Function(stmt) => Some(&stmt.parameters),
@@ -149,19 +153,24 @@ impl FunctionLike<'_> {
     /// function from `operator` module.
     fn try_fix(
         &self,
-        operator: &'static str,
+        operator: &Operator,
         importer: &Importer,
         semantic: &SemanticModel,
     ) -> Result<Option<Fix>> {
         match self {
             Self::Lambda(_) => {
                 let (edit, binding) = importer.get_or_import_symbol(
-                    &ImportRequest::import("operator", operator),
+                    &ImportRequest::import("operator", operator.name),
                     self.start(),
                     semantic,
                 )?;
+                let content = if operator.args.is_empty() {
+                    binding
+                } else {
+                    format!("{binding}({})", operator.args.join(", "))
+                };
                 Ok(Some(Fix::safe_edits(
-                    Edit::range_replacement(binding, self.range()),
+                    Edit::range_replacement(content, self.range()),
                     [edit],
                 )))
             }
@@ -170,12 +179,112 @@ impl FunctionLike<'_> {
     }
 }
 
-/// Return the name of the `operator` implemented by the given expression.
-fn get_operator(expr: &Expr, params: &ast::Parameters) -> Option<&'static str> {
+/// Convert the slice expression to the string representation of `slice` call.
+/// For example, expression `1:2` will be `slice(1, 2)`, and `:` will be `slice(None)`.
+fn slice_expr_to_slice_call(slice: &ExprSlice, locator: &Locator) -> String {
+    let stringify = |expr: Option<&Expr>| expr.map_or("None", |expr| locator.slice(expr));
+    match (
+        slice.lower.as_deref(),
+        slice.upper.as_deref(),
+        slice.step.as_deref(),
+    ) {
+        (lower, upper, step @ Some(_)) => format!(
+            "slice({}, {}, {})",
+            stringify(lower),
+            stringify(upper),
+            stringify(step)
+        ),
+        (None, upper, None) => format!("slice({})", stringify(upper)),
+        (lower @ Some(_), upper, None) => {
+            format!("slice({}, {})", stringify(lower), stringify(upper))
+        }
+    }
+}
+
+/// Convert the given expression to a string representation, suitable to be a function argument.
+fn subscript_slice_to_string<'a>(expr: &Expr, locator: &Locator<'a>) -> Cow<'a, str> {
+    if let Expr::Slice(expr_slice) = expr {
+        Cow::Owned(slice_expr_to_slice_call(expr_slice, locator))
+    } else {
+        Cow::Borrowed(locator.slice(expr))
+    }
+}
+
+/// Return the `operator` implemented by given subscript expression.
+fn itemgetter_op(expr: &ExprSubscript, params: &Parameters, locator: &Locator) -> Option<Operator> {
+    let [arg] = params.args.as_slice() else {
+        return None;
+    };
+    if !is_same_expression(arg, &expr.value) {
+        return None;
+    };
+    Some(Operator {
+        name: "itemgetter",
+        args: vec![subscript_slice_to_string(expr.slice.as_ref(), locator).to_string()],
+    })
+}
+
+/// Return the `operator` implemented by given tuple expression.
+fn itemgetter_op_tuple(
+    expr: &ExprTuple,
+    params: &Parameters,
+    locator: &Locator,
+) -> Option<Operator> {
+    let [arg] = params.args.as_slice() else {
+        return None;
+    };
+    if expr.elts.is_empty() {
+        return None;
+    }
+    Some(Operator {
+        name: "itemgetter",
+        args: expr
+            .elts
+            .iter()
+            .map(|expr| {
+                expr.as_subscript_expr()
+                    .filter(|expr| is_same_expression(arg, &expr.value))
+                    .map(|expr| expr.slice.as_ref())
+                    .map(|slice| subscript_slice_to_string(slice, locator).to_string())
+            })
+            .collect::<Option<Vec<_>>>()?,
+    })
+}
+
+#[derive(Eq, PartialEq, Debug)]
+struct Operator {
+    name: &'static str,
+    args: Vec<String>,
+}
+
+impl From<&'static str> for Operator {
+    fn from(value: &'static str) -> Self {
+        Self {
+            name: value,
+            args: vec![],
+        }
+    }
+}
+
+impl Display for Operator {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.name)?;
+        if self.args.is_empty() {
+            Ok(())
+        } else {
+            write!(f, "({})", self.args.join(", "))
+        }
+    }
+}
+
+/// Return the `operator` implemented by the given expression.
+fn get_operator(expr: &Expr, params: &Parameters, locator: &Locator) -> Option<Operator> {
     match expr {
-        Expr::UnaryOp(expr) => unary_op(expr, params),
-        Expr::BinOp(expr) => bin_op(expr, params),
-        Expr::Compare(expr) => cmp_op(expr, params),
+        Expr::UnaryOp(expr) => unary_op(expr, params).map(Operator::from),
+        Expr::BinOp(expr) => bin_op(expr, params).map(Operator::from),
+        Expr::Compare(expr) => cmp_op(expr, params).map(Operator::from),
+        Expr::Subscript(expr) => itemgetter_op(expr, params, locator),
+        Expr::Tuple(expr) => itemgetter_op_tuple(expr, params, locator),
         _ => None,
     }
 }
@@ -187,7 +296,7 @@ enum FunctionLikeKind {
 }
 
 /// Return the name of the `operator` implemented by the given unary expression.
-fn unary_op(expr: &ast::ExprUnaryOp, params: &ast::Parameters) -> Option<&'static str> {
+fn unary_op(expr: &ast::ExprUnaryOp, params: &Parameters) -> Option<&'static str> {
     let [arg] = params.args.as_slice() else {
         return None;
     };
@@ -203,7 +312,7 @@ fn unary_op(expr: &ast::ExprUnaryOp, params: &ast::Parameters) -> Option<&'stati
 }
 
 /// Return the name of the `operator` implemented by the given binary expression.
-fn bin_op(expr: &ast::ExprBinOp, params: &ast::Parameters) -> Option<&'static str> {
+fn bin_op(expr: &ast::ExprBinOp, params: &Parameters) -> Option<&'static str> {
     let [arg1, arg2] = params.args.as_slice() else {
         return None;
     };
@@ -228,7 +337,7 @@ fn bin_op(expr: &ast::ExprBinOp, params: &ast::Parameters) -> Option<&'static st
 }
 
 /// Return the name of the `operator` implemented by the given comparison expression.
-fn cmp_op(expr: &ast::ExprCompare, params: &ast::Parameters) -> Option<&'static str> {
+fn cmp_op(expr: &ast::ExprCompare, params: &Parameters) -> Option<&'static str> {
     let [arg1, arg2] = params.args.as_slice() else {
         return None;
     };
@@ -240,71 +349,19 @@ fn cmp_op(expr: &ast::ExprCompare, params: &ast::Parameters) -> Option<&'static
     };
 
     match op {
-        ast::CmpOp::Eq => {
-            if match_arguments(arg1, arg2, &expr.left, right) {
-                Some("eq")
-            } else {
-                None
-            }
-        }
-        ast::CmpOp::NotEq => {
-            if match_arguments(arg1, arg2, &expr.left, right) {
-                Some("ne")
-            } else {
-                None
-            }
-        }
-        ast::CmpOp::Lt => {
-            if match_arguments(arg1, arg2, &expr.left, right) {
-                Some("lt")
-            } else {
-                None
-            }
-        }
-        ast::CmpOp::LtE => {
-            if match_arguments(arg1, arg2, &expr.left, right) {
-                Some("le")
-            } else {
-                None
-            }
-        }
-        ast::CmpOp::Gt => {
-            if match_arguments(arg1, arg2, &expr.left, right) {
-                Some("gt")
-            } else {
-                None
-            }
-        }
-        ast::CmpOp::GtE => {
-            if match_arguments(arg1, arg2, &expr.left, right) {
-                Some("ge")
-            } else {
-                None
-            }
-        }
-        ast::CmpOp::Is => {
-            if match_arguments(arg1, arg2, &expr.left, right) {
-                Some("is_")
-            } else {
-                None
-            }
-        }
-        ast::CmpOp::IsNot => {
-            if match_arguments(arg1, arg2, &expr.left, right) {
-                Some("is_not")
-            } else {
-                None
-            }
-        }
+        ast::CmpOp::Eq => match_arguments(arg1, arg2, &expr.left, right).then_some("eq"),
+        ast::CmpOp::NotEq => match_arguments(arg1, arg2, &expr.left, right).then_some("ne"),
+        ast::CmpOp::Lt => match_arguments(arg1, arg2, &expr.left, right).then_some("lt"),
+        ast::CmpOp::LtE => match_arguments(arg1, arg2, &expr.left, right).then_some("le"),
+        ast::CmpOp::Gt => match_arguments(arg1, arg2, &expr.left, right).then_some("gt"),
+        ast::CmpOp::GtE => match_arguments(arg1, arg2, &expr.left, right).then_some("ge"),
+        ast::CmpOp::Is => match_arguments(arg1, arg2, &expr.left, right).then_some("is_"),
+        ast::CmpOp::IsNot => match_arguments(arg1, arg2, &expr.left, right).then_some("is_not"),
         ast::CmpOp::In => {
             // Note: `operator.contains` reverses the order of arguments. That is:
             // `operator.contains` is equivalent to `lambda x, y: y in x`, rather than
             // `lambda x, y: x in y`.
-            if match_arguments(arg1, arg2, right, &expr.left) {
-                Some("contains")
-            } else {
-                None
-            }
+            match_arguments(arg1, arg2, right, &expr.left).then_some("contains")
         }
         ast::CmpOp::NotIn => None,
     }