Skip to content

Commit

Permalink
Minor tweaks
Browse files Browse the repository at this point in the history
  • Loading branch information
charliermarsh committed Apr 7, 2024
1 parent a444f30 commit cd6afb8
Show file tree
Hide file tree
Showing 3 changed files with 72 additions and 79 deletions.
9 changes: 6 additions & 3 deletions crates/ruff_linter/resources/test/fixtures/refurb/FURB118.py
Original file line number Diff line number Diff line change
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 Down Expand Up @@ -45,8 +45,9 @@ 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
Expand All @@ -63,9 +64,11 @@ def add(x, y):
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
114 changes: 53 additions & 61 deletions crates/ruff_linter/src/rules/refurb/rules/reimplemented_operator.rs
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
use std::borrow::Cow;
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, 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 @@ -74,7 +75,7 @@ pub(crate) fn reimplemented_operator(checker: &mut Checker, target: &FunctionLik
return;
};
let Some(body) = target.body() else { return };
let Some(operator) = get_operator(checker, body, params) else {
let Some(operator) = get_operator(body, params, checker.locator()) else {
return;
};
let fix = target.try_fix(&operator, checker.importer(), checker.semantic());
Expand Down Expand Up @@ -118,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 @@ -163,10 +164,10 @@ impl FunctionLike<'_> {
self.start(),
semantic,
)?;
let content = if let Some(args) = operator.args.as_ref() {
format!("{binding}({args})")
} else {
let content = if operator.args.is_empty() {
binding
} else {
format!("{binding}({})", operator.args.join(", "))
};
Ok(Some(Fix::safe_edits(
Edit::range_replacement(content, self.range()),
Expand All @@ -180,40 +181,37 @@ impl FunctionLike<'_> {

/// 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));
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 (
expr_slice.lower.as_ref(),
expr_slice.upper.as_ref(),
expr_slice.step.as_ref(),
slice.lower.as_deref(),
slice.upper.as_deref(),
slice.step.as_deref(),
) {
(l, u, s @ Some(_)) => format!(
(lower, upper, step @ Some(_)) => format!(
"slice({}, {}, {})",
stringify(l),
stringify(u),
stringify(s)
stringify(lower),
stringify(upper),
stringify(step)
),
(None, u, None) => format!("slice({})", stringify(u)),
(l @ Some(_), u, None) => format!("slice({}, {})", stringify(l), stringify(u)),
(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(checker: &mut Checker, expr: &Expr) -> String {
fn subscript_slice_to_string<'a>(expr: &Expr, locator: &Locator<'a>) -> Cow<'a, str> {
if let Expr::Slice(expr_slice) = expr {
slice_expr_to_slice_call(checker, expr_slice)
Cow::Owned(slice_expr_to_slice_call(expr_slice, locator))
} else {
checker.generator().expr(expr)
Cow::Borrowed(locator.slice(expr))
}
}

/// Return the `operator` implemented by given subscript expression.
fn itemgetter_op(
checker: &mut Checker,
expr: &ExprSubscript,
params: &Parameters,
) -> Option<Operator> {
fn itemgetter_op(expr: &ExprSubscript, params: &Parameters, locator: &Locator) -> Option<Operator> {
let [arg] = params.args.as_slice() else {
return None;
};
Expand All @@ -222,77 +220,71 @@ fn itemgetter_op(
};
Some(Operator {
name: "itemgetter",
args: Some(subscript_slice_to_string(checker, expr.slice.as_ref())),
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(
checker: &mut Checker,
expr: &ExprTuple,
params: &Parameters,
locator: &Locator,
) -> 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(", "),
),
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: Option<String>,
args: Vec<String>,
}

impl From<&'static str> for Operator {
fn from(value: &'static str) -> Self {
Self {
name: value,
args: None,
args: vec![],
}
}
}

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})"))
if self.args.is_empty() {
Ok(())
} else {
write!(f, "({})", self.args.join(", "))
}
}
}

/// Return the `operator` implemented by the given expression.
fn get_operator(checker: &mut Checker, expr: &Expr, params: &ast::Parameters) -> Option<Operator> {
fn get_operator(expr: &Expr, params: &Parameters, locator: &Locator) -> Option<Operator> {
match expr {
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),
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 @@ -304,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 @@ -320,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 @@ -345,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 Down
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,7 @@ FURB118.py:10:14: FURB118 [*] Use `operator.matmul` instead of defining a lambda
11 |+op_matmutl = operator.matmul
11 12 | op_truediv = lambda x, y: x / y
12 13 | op_mod = lambda x, y: x % y
13 14 | op_pow = lambda x, y: x ** y
13 14 | op_pow = lambda x, y: x**y

FURB118.py:11:14: FURB118 [*] Use `operator.truediv` instead of defining a lambda
|
Expand All @@ -196,7 +196,7 @@ FURB118.py:11:14: FURB118 [*] Use `operator.truediv` instead of defining a lambd
11 | op_truediv = lambda x, y: x / y
| ^^^^^^^^^^^^^^^^^^ FURB118
12 | op_mod = lambda x, y: x % y
13 | op_pow = lambda x, y: x ** y
13 | op_pow = lambda x, y: x**y
|
= help: Replace with `operator.truediv`

Expand All @@ -213,7 +213,7 @@ FURB118.py:11:14: FURB118 [*] Use `operator.truediv` instead of defining a lambd
11 |-op_truediv = lambda x, y: x / y
12 |+op_truediv = operator.truediv
12 13 | op_mod = lambda x, y: x % y
13 14 | op_pow = lambda x, y: x ** y
13 14 | op_pow = lambda x, y: x**y
14 15 | op_lshift = lambda x, y: x << y

FURB118.py:12:10: FURB118 [*] Use `operator.mod` instead of defining a lambda
Expand All @@ -222,7 +222,7 @@ FURB118.py:12:10: FURB118 [*] Use `operator.mod` instead of defining a lambda
11 | op_truediv = lambda x, y: x / y
12 | op_mod = lambda x, y: x % y
| ^^^^^^^^^^^^^^^^^^ FURB118
13 | op_pow = lambda x, y: x ** y
13 | op_pow = lambda x, y: x**y
14 | op_lshift = lambda x, y: x << y
|
= help: Replace with `operator.mod`
Expand All @@ -239,16 +239,16 @@ FURB118.py:12:10: FURB118 [*] Use `operator.mod` instead of defining a lambda
11 12 | op_truediv = lambda x, y: x / y
12 |-op_mod = lambda x, y: x % y
13 |+op_mod = operator.mod
13 14 | op_pow = lambda x, y: x ** y
13 14 | op_pow = lambda x, y: x**y
14 15 | op_lshift = lambda x, y: x << y
15 16 | op_rshift = lambda x, y: x >> y

FURB118.py:13:10: FURB118 [*] Use `operator.pow` instead of defining a lambda
|
11 | op_truediv = lambda x, y: x / y
12 | op_mod = lambda x, y: x % y
13 | op_pow = lambda x, y: x ** y
| ^^^^^^^^^^^^^^^^^^^ FURB118
13 | op_pow = lambda x, y: x**y
| ^^^^^^^^^^^^^^^^^ FURB118
14 | op_lshift = lambda x, y: x << y
15 | op_rshift = lambda x, y: x >> y
|
Expand All @@ -264,7 +264,7 @@ FURB118.py:13:10: FURB118 [*] Use `operator.pow` instead of defining a lambda
10 11 | op_matmutl = lambda x, y: x @ y
11 12 | op_truediv = lambda x, y: x / y
12 13 | op_mod = lambda x, y: x % y
13 |-op_pow = lambda x, y: x ** y
13 |-op_pow = lambda x, y: x**y
14 |+op_pow = operator.pow
14 15 | op_lshift = lambda x, y: x << y
15 16 | op_rshift = lambda x, y: x >> y
Expand All @@ -273,7 +273,7 @@ FURB118.py:13:10: FURB118 [*] Use `operator.pow` instead of defining a lambda
FURB118.py:14:13: FURB118 [*] Use `operator.lshift` instead of defining a lambda
|
12 | op_mod = lambda x, y: x % y
13 | op_pow = lambda x, y: x ** y
13 | op_pow = lambda x, y: x**y
14 | op_lshift = lambda x, y: x << y
| ^^^^^^^^^^^^^^^^^^^ FURB118
15 | op_rshift = lambda x, y: x >> y
Expand All @@ -290,7 +290,7 @@ FURB118.py:14:13: FURB118 [*] Use `operator.lshift` instead of defining a lambda
--------------------------------------------------------------------------------
11 12 | op_truediv = lambda x, y: x / y
12 13 | op_mod = lambda x, y: x % y
13 14 | op_pow = lambda x, y: x ** y
13 14 | op_pow = lambda x, y: x**y
14 |-op_lshift = lambda x, y: x << y
15 |+op_lshift = operator.lshift
15 16 | op_rshift = lambda x, y: x >> y
Expand All @@ -299,7 +299,7 @@ FURB118.py:14:13: FURB118 [*] Use `operator.lshift` instead of defining a lambda

FURB118.py:15:13: FURB118 [*] Use `operator.rshift` instead of defining a lambda
|
13 | op_pow = lambda x, y: x ** y
13 | op_pow = lambda x, y: x**y
14 | op_lshift = lambda x, y: x << y
15 | op_rshift = lambda x, y: x >> y
| ^^^^^^^^^^^^^^^^^^^ FURB118
Expand All @@ -316,7 +316,7 @@ FURB118.py:15:13: FURB118 [*] Use `operator.rshift` instead of defining a lambda
4 5 | op_pos = lambda x: +x
--------------------------------------------------------------------------------
12 13 | op_mod = lambda x, y: x % y
13 14 | op_pow = lambda x, y: x ** y
13 14 | op_pow = lambda x, y: x**y
14 15 | op_lshift = lambda x, y: x << y
15 |-op_rshift = lambda x, y: x >> y
16 |+op_rshift = operator.rshift
Expand All @@ -342,7 +342,7 @@ FURB118.py:16:12: FURB118 [*] Use `operator.or_` instead of defining a lambda
3 4 | op_not = lambda x: not x
4 5 | op_pos = lambda x: +x
--------------------------------------------------------------------------------
13 14 | op_pow = lambda x, y: x ** y
13 14 | op_pow = lambda x, y: x**y
14 15 | op_lshift = lambda x, y: x << y
15 16 | op_rshift = lambda x, y: x >> y
16 |-op_bitor = lambda x, y: x | y
Expand Down Expand Up @@ -801,7 +801,5 @@ FURB118.py:45:5: FURB118 Use `operator.add` instead of defining a function
| _____^
46 | | return x + y
| |____________________^ FURB118
47 |
48 | # OK.
|
= help: Replace with `operator.add`

0 comments on commit cd6afb8

Please sign in to comment.