Skip to content

Commit

Permalink
Implement reimplemented_operator (FURB118) (#9171)
Browse files Browse the repository at this point in the history
## Summary

Implement
[FURB118](https://github.com/dosisod/refurb/blob/master/docs/checks.md#furb118-use-operator)
that recommends, for example, that `lambda x, y: x + y` is replaced with
`operator.add`. Part of #1348.

## Test Plan

Added test cases.
  • Loading branch information
siiptuo authored Dec 18, 2023
1 parent 0977fa9 commit c532089
Show file tree
Hide file tree
Showing 9 changed files with 1,033 additions and 1 deletion.
56 changes: 56 additions & 0 deletions crates/ruff_linter/resources/test/fixtures/refurb/FURB118.py
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -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());
}
}
}
}
3 changes: 3 additions & 0 deletions crates/ruff_linter/src/checkers/ast/analyze/statement.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
1 change: 1 addition & 0 deletions crates/ruff_linter/src/codes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down
1 change: 1 addition & 0 deletions crates/ruff_linter/src/rules/refurb/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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"))]
Expand Down
2 changes: 2 additions & 0 deletions crates/ruff_linter/src/rules/refurb/rules/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::*;
Expand All @@ -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;
Expand Down
254 changes: 254 additions & 0 deletions crates/ruff_linter/src/rules/refurb/rules/reimplemented_operator.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,254 @@
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_python_semantic::SemanticModel;
use ruff_text_size::{Ranged, TextRange};

use crate::checkers::ast::Checker;
use crate::importer::{ImportRequest, Importer};

/// ## What it does
/// Checks for lambda expressions and function definitions that can be replaced
/// with a function from the `operator` module.
///
/// ## Why is this bad?
/// The `operator` module provides functions that implement the same functionality
/// as the corresponding operators. For example, `operator.add` is equivalent to
/// `lambda x, y: x + y`. Using the functions from the `operator` module is more
/// concise and communicates the intent of the code more clearly.
///
/// ## 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: &FunctionLike) {
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(operator, checker.importer(), checker.semantic()));
checker.diagnostics.push(diagnostic);
}

/// Candidate for lambda expression or function definition consisting of a return statement.
#[derive(Debug)]
pub(crate) enum FunctionLike<'a> {
Lambda(&'a ast::ExprLambda),
Function(&'a ast::StmtFunctionDef),
}

impl<'a> From<&'a ast::ExprLambda> for FunctionLike<'a> {
fn from(lambda: &'a ast::ExprLambda) -> Self {
Self::Lambda(lambda)
}
}

impl<'a> From<&'a ast::StmtFunctionDef> for FunctionLike<'a> {
fn from(function: &'a ast::StmtFunctionDef) -> Self {
Self::Function(function)
}
}

impl Ranged for FunctionLike<'_> {
fn range(&self) -> TextRange {
match self {
Self::Lambda(expr) => expr.range(),
Self::Function(stmt) => stmt.range(),
}
}
}

impl FunctionLike<'_> {
/// Return the [`ast::Parameters`] of the function-like node.
fn parameters(&self) -> Option<&ast::Parameters> {
match self {
Self::Lambda(expr) => expr.parameters.as_deref(),
Self::Function(stmt) => Some(&stmt.parameters),
}
}

/// Return the body of the function-like node.
///
/// If the node is a function definition that consists of more than a single return statement,
/// returns `None`.
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,
},
}
}

/// Return the display kind of the function-like node.
fn kind(&self) -> &'static str {
match self {
Self::Lambda(_) => "lambda",
Self::Function(_) => "function",
}
}

/// Attempt to fix the function-like node by replacing it with a call to the corresponding
/// function from `operator` module.
fn try_fix(
&self,
operator: &'static str,
importer: &Importer,
semantic: &SemanticModel,
) -> Result<Fix> {
match self {
Self::Lambda(_) => {
let (edit, binding) = importer.get_or_import_symbol(
&ImportRequest::import("operator", operator),
self.start(),
semantic,
)?;
Ok(Fix::safe_edits(
Edit::range_replacement(binding, self.range()),
[edit],
))
}
Self::Function(_) => bail!("No fix available"),
}
}
}

/// Return the name of the `operator` implemented by the given expression.
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,
}
}

/// Return the name of the `operator` implemented by the given unary expression.
fn unary_op(expr: &ast::ExprUnaryOp, params: &ast::Parameters) -> Option<&'static str> {
let [arg] = params.args.as_slice() else {
return None;
};
if !is_same_expression(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",
})
}

/// Return the name of the `operator` implemented by the given binary expression.
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_expression(arg1, &expr.left) || !is_same_expression(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",
})
}

/// Return the name of the `operator` implemented by the given comparison expression.
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_expression(arg1, &expr.left) || !is_same_expression(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,
}
}

/// Returns `true` if the given argument is the "same" as the given expression. For example, if
/// the argument has a default, it is not considered the same as any expression; if both match the
/// same name, they are considered the same.
fn is_same_expression(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
}
}
Loading

0 comments on commit c532089

Please sign in to comment.