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

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

@@ -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):
@@ -48,6 +52,12 @@ def add(x, y):
 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):

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

@@ -1,8 +1,11 @@
+use std::fmt::{Debug, Display, Formatter};
+
 use anyhow::Result;
+use itertools::Itertools;
 
 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_text_size::{Ranged, TextRange};
 
@@ -39,7 +42,7 @@ use crate::importer::{ImportRequest, Importer};
 /// ## References
 #[violation]
 pub struct ReimplementedOperator {
-    operator: &'static str,
+    operator: Operator,
     target: FunctionLikeKind,
 }
 
@@ -71,18 +74,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(checker, body, params) 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);
 }
 
@@ -149,19 +152,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 let Some(args) = operator.args.as_ref() {
+                    format!("{binding}({args})")
+                } else {
+                    binding
+                };
                 Ok(Some(Fix::safe_edits(
-                    Edit::range_replacement(binding, self.range()),
+                    Edit::range_replacement(content, self.range()),
                     [edit],
                 )))
             }
@@ -170,12 +178,121 @@ 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(checker: &mut Checker, expr_slice: &ExprSlice) -> String {
+    let stringify =
+        |x: Option<&Box<Expr>>| x.map_or("None".into(), |x| checker.generator().expr(x));
+    match (
+        expr_slice.lower.as_ref(),
+        expr_slice.upper.as_ref(),
+        expr_slice.step.as_ref(),
+    ) {
+        (l, u, s @ Some(_)) => format!(
+            "slice({}, {}, {})",
+            stringify(l),
+            stringify(u),
+            stringify(s)
+        ),
+        (None, u, None) => format!("slice({})", stringify(u)),
+        (l @ Some(_), u, None) => format!("slice({}, {})", stringify(l), stringify(u)),
+    }
+}
+
+/// Convert the given expression to a string representation, suitable to be a function argument.
+fn subscript_slice_to_string(checker: &mut Checker, expr: &Expr) -> String {
+    if let Expr::Slice(expr_slice) = expr {
+        slice_expr_to_slice_call(checker, expr_slice)
+    } else {
+        checker.generator().expr(expr)
+    }
+}
+
+/// Return the `operator` implemented by given subscript expression.
+fn itemgetter_op(
+    checker: &mut Checker,
+    expr: &ExprSubscript,
+    params: &Parameters,
+) -> 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: Some(subscript_slice_to_string(checker, expr.slice.as_ref())),
+    })
+}
+
+/// Return the `operator` implemented by given tuple expression.
+fn itemgetter_op_tuple(
+    checker: &mut Checker,
+    expr: &ExprTuple,
+    params: &Parameters,
+) -> Option<Operator> {
+    let [arg] = params.args.as_slice() else {
+        return None;
+    };
+    if expr.elts.is_empty() {
+        return None;
+    }
+    if !expr.elts.iter().all(|expr| {
+        expr.as_subscript_expr()
+            .is_some_and(|expr| is_same_expression(arg, &expr.value))
+    }) {
+        return None;
+    }
+    Some(Operator {
+        name: "itemgetter",
+        args: Some(
+            expr.elts
+                .iter()
+                .map(|expr| {
+                    subscript_slice_to_string(
+                        checker,
+                        // unwrap is safe, because we check that all elts are subscripts
+                        expr.as_subscript_expr().unwrap().slice.as_ref(),
+                    )
+                })
+                .join(", "),
+        ),
+    })
+}
+
+#[derive(Eq, PartialEq, Debug)]
+struct Operator {
+    name: &'static str,
+    args: Option<String>,
+}
+
+impl From<&'static str> for Operator {
+    fn from(value: &'static str) -> Self {
+        Self {
+            name: value,
+            args: None,
+        }
+    }
+}
+
+impl Display for Operator {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.name)?;
+        self.args
+            .as_ref()
+            .map_or(Ok(()), |args| write!(f, "({args})"))
+    }
+}
+
+/// Return the `operator` implemented by the given expression.
+fn get_operator(checker: &mut Checker, expr: &Expr, params: &ast::Parameters) -> 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(Into::into),
+        Expr::BinOp(expr) => bin_op(expr, params).map(Into::into),
+        Expr::Compare(expr) => cmp_op(expr, params).map(Into::into),
+        Expr::Subscript(expr) => itemgetter_op(checker, expr, params),
+        Expr::Tuple(expr) => itemgetter_op_tuple(checker, expr, params),
         _ => None,
     }
 }
@@ -240,71 +357,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,
     }