Skip to content

Commit

Permalink
[refurb] Support itemgetter in reimplemented-operator (`FURB118…
Browse files Browse the repository at this point in the history
…`) (#10526)

## Summary
Lint about function like expressions which are equivalent to
`operator.itemgetter`.
See:
#1348 (comment)

## Test Plan
cargo test
  • Loading branch information
alex-700 committed Apr 7, 2024
1 parent 2a51dcf commit 6050bab
Show file tree
Hide file tree
Showing 3 changed files with 289 additions and 115 deletions.
19 changes: 16 additions & 3 deletions crates/ruff_linter/resources/test/fixtures/refurb/FURB118.py
Expand Up @@ -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
Expand All @@ -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):
Expand All @@ -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
215 changes: 136 additions & 79 deletions 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;
Expand Down Expand Up @@ -39,7 +43,7 @@ use crate::importer::{ImportRequest, Importer};
/// ## References
#[violation]
pub struct ReimplementedOperator {
operator: &'static str,
operator: Operator,
target: FunctionLikeKind,
}

Expand Down Expand Up @@ -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);
}

Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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],
)))
}
Expand All @@ -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,
}
}
Expand All @@ -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;
};
Expand All @@ -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;
};
Expand All @@ -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;
};
Expand All @@ -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,
}
Expand Down

0 comments on commit 6050bab

Please sign in to comment.