Skip to content

Commit

Permalink
[pylint] Implement if-stmt-min-max (PLR1730, PLR1731) (#10002)
Browse files Browse the repository at this point in the history
  • Loading branch information
tibor-reiss committed Apr 6, 2024
1 parent ee4bff3 commit 3194f90
Show file tree
Hide file tree
Showing 8 changed files with 636 additions and 0 deletions.
136 changes: 136 additions & 0 deletions crates/ruff_linter/resources/test/fixtures/pylint/if_stmt_min_max.py
@@ -0,0 +1,136 @@
# pylint: disable=missing-docstring, invalid-name, too-few-public-methods, redefined-outer-name

value = 10
value2 = 0
value3 = 3

# Positive
if value < 10: # [max-instead-of-if]
value = 10

if value <= 10: # [max-instead-of-if]
value = 10

if value < value2: # [max-instead-of-if]
value = value2

if value > 10: # [min-instead-of-if]
value = 10

if value >= 10: # [min-instead-of-if]
value = 10

if value > value2: # [min-instead-of-if]
value = value2


class A:
def __init__(self):
self.value = 13


A1 = A()
if A1.value < 10: # [max-instead-of-if]
A1.value = 10

if A1.value > 10: # [min-instead-of-if]
A1.value = 10


class AA:
def __init__(self, value):
self.value = value

def __gt__(self, b):
return self.value > b

def __ge__(self, b):
return self.value >= b

def __lt__(self, b):
return self.value < b

def __le__(self, b):
return self.value <= b


A1 = AA(0)
A2 = AA(3)

if A2 < A1: # [max-instead-of-if]
A2 = A1

if A2 <= A1: # [max-instead-of-if]
A2 = A1

if A2 > A1: # [min-instead-of-if]
A2 = A1

if A2 >= A1: # [min-instead-of-if]
A2 = A1

# Negative
if value < 10:
value = 2

if value <= 3:
value = 5

if value < 10:
value = 2
value2 = 3

if value < value2:
value = value3

if value < 5:
value = value3

if 2 < value <= 3:
value = 1

if value < 10:
value = 10
else:
value = 3

if value <= 3:
value = 5
elif value == 3:
value = 2

if value > 10:
value = 2

if value >= 3:
value = 5

if value > 10:
value = 2
value2 = 3

if value > value2:
value = value3

if value > 5:
value = value3

if 2 > value >= 3:
value = 1

if value > 10:
value = 10
else:
value = 3

if value >= 3:
value = 5
elif value == 3:
value = 2

# Parenthesized expressions
if value.attr > 3:
(
value.
attr
) = 3
3 changes: 3 additions & 0 deletions crates/ruff_linter/src/checkers/ast/analyze/statement.rs
Expand Up @@ -1117,6 +1117,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
if checker.enabled(Rule::TooManyBooleanExpressions) {
pylint::rules::too_many_boolean_expressions(checker, if_);
}
if checker.enabled(Rule::IfStmtMinMax) {
pylint::rules::if_stmt_min_max(checker, if_);
}
if checker.source_type.is_stub() {
if checker.any_enabled(&[
Rule::UnrecognizedVersionInfoCheck,
Expand Down
1 change: 1 addition & 0 deletions crates/ruff_linter/src/codes.rs
Expand Up @@ -287,6 +287,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
(Pylint, "R1711") => (RuleGroup::Stable, rules::pylint::rules::UselessReturn),
(Pylint, "R1714") => (RuleGroup::Stable, rules::pylint::rules::RepeatedEqualityComparison),
(Pylint, "R1722") => (RuleGroup::Stable, rules::pylint::rules::SysExitAlias),
(Pylint, "R1730") => (RuleGroup::Preview, rules::pylint::rules::IfStmtMinMax),
(Pylint, "R1733") => (RuleGroup::Preview, rules::pylint::rules::UnnecessaryDictIndexLookup),
(Pylint, "R1736") => (RuleGroup::Preview, rules::pylint::rules::UnnecessaryListIndexLookup),
(Pylint, "R2004") => (RuleGroup::Stable, rules::pylint::rules::MagicValueComparison),
Expand Down
1 change: 1 addition & 0 deletions crates/ruff_linter/src/rules/pylint/mod.rs
Expand Up @@ -47,6 +47,7 @@ mod tests {
#[test_case(Rule::EqWithoutHash, Path::new("eq_without_hash.py"))]
#[test_case(Rule::EmptyComment, Path::new("empty_comment.py"))]
#[test_case(Rule::ManualFromImport, Path::new("import_aliasing.py"))]
#[test_case(Rule::IfStmtMinMax, Path::new("if_stmt_min_max.py"))]
#[test_case(Rule::SingleStringSlots, Path::new("single_string_slots.py"))]
#[test_case(Rule::SysExitAlias, Path::new("sys_exit_alias_0.py"))]
#[test_case(Rule::SysExitAlias, Path::new("sys_exit_alias_1.py"))]
Expand Down
196 changes: 196 additions & 0 deletions crates/ruff_linter/src/rules/pylint/rules/if_stmt_min_max.rs
@@ -0,0 +1,196 @@
use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::comparable::ComparableExpr;
use ruff_python_ast::parenthesize::parenthesized_range;
use ruff_python_ast::{self as ast, CmpOp, Stmt};
use ruff_text_size::Ranged;

use crate::checkers::ast::Checker;
use crate::fix::snippet::SourceCodeSnippet;

/// ## What it does
/// Checks for `if` statements that can be replaced with `min()` or `max()`
/// calls.
///
/// ## Why is this bad?
/// An `if` statement that selects the lesser or greater of two sub-expressions
/// can be replaced with a `min()` or `max()` call respectively. When possible,
/// prefer `min()` and `max()`, as they're more concise and readable than the
/// equivalent `if` statements.
///
/// ## Example
/// ```python
/// if score > highest_score:
/// highest_score = score
/// ```
///
/// Use instead:
/// ```python
/// highest_score = max(highest_score, score)
/// ```
///
/// ## References
/// - [Python documentation: max function](https://docs.python.org/3/library/functions.html#max)
/// - [Python documentation: min function](https://docs.python.org/3/library/functions.html#min)
#[violation]
pub struct IfStmtMinMax {
min_max: MinMax,
replacement: SourceCodeSnippet,
}

impl Violation for IfStmtMinMax {
const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;

#[derive_message_formats]
fn message(&self) -> String {
let Self {
min_max,
replacement,
} = self;
if let Some(replacement) = replacement.full_display() {
format!("Replace `if` statement with `{replacement}`")
} else {
format!("Replace `if` statement with `{min_max}` call")
}
}

fn fix_title(&self) -> Option<String> {
let Self {
min_max,
replacement,
} = self;
if let Some(replacement) = replacement.full_display() {
Some(format!("Replace with `{replacement}`"))
} else {
Some(format!("Replace with `{min_max}` call"))
}
}
}

/// R1730, R1731
pub(crate) fn if_stmt_min_max(checker: &mut Checker, stmt_if: &ast::StmtIf) {
let ast::StmtIf {
test,
body,
elif_else_clauses,
range: _,
} = stmt_if;

if !elif_else_clauses.is_empty() {
return;
}

let [body @ Stmt::Assign(ast::StmtAssign {
targets: body_targets,
value: body_value,
..
})] = body.as_slice()
else {
return;
};
let [body_target] = body_targets.as_slice() else {
return;
};

let Some(ast::ExprCompare {
ops,
left,
comparators,
..
}) = test.as_compare_expr()
else {
return;
};

// Ignore, e.g., `foo < bar < baz`.
let [op] = &**ops else {
return;
};

// Determine whether to use `min()` or `max()`, and whether to flip the
// order of the arguments, which is relevant for breaking ties.
let (min_max, flip_args) = match op {
CmpOp::Gt => (MinMax::Max, true),
CmpOp::GtE => (MinMax::Max, false),
CmpOp::Lt => (MinMax::Min, true),
CmpOp::LtE => (MinMax::Min, false),
_ => return,
};

let [right] = &**comparators else {
return;
};

let _min_or_max = match op {
CmpOp::Gt | CmpOp::GtE => MinMax::Min,
CmpOp::Lt | CmpOp::LtE => MinMax::Max,
_ => return,
};

let left_cmp = ComparableExpr::from(left);
let body_target_cmp = ComparableExpr::from(body_target);
let right_statement_cmp = ComparableExpr::from(right);
let body_value_cmp = ComparableExpr::from(body_value);
if left_cmp != body_target_cmp || right_statement_cmp != body_value_cmp {
return;
}

let (arg1, arg2) = if flip_args {
(left.as_ref(), right)
} else {
(right, left.as_ref())
};

let replacement = format!(
"{} = {min_max}({}, {})",
checker.locator().slice(
parenthesized_range(
body_target.into(),
body.into(),
checker.indexer().comment_ranges(),
checker.locator().contents()
)
.unwrap_or(body_target.range())
),
checker.locator().slice(arg1),
checker.locator().slice(arg2),
);

let mut diagnostic = Diagnostic::new(
IfStmtMinMax {
min_max,
replacement: SourceCodeSnippet::from_str(replacement.as_str()),
},
stmt_if.range(),
);

if checker.semantic().is_builtin(min_max.as_str()) {
diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
replacement,
stmt_if.range(),
)));
}

checker.diagnostics.push(diagnostic);
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
enum MinMax {
Min,
Max,
}

impl MinMax {
fn as_str(self) -> &'static str {
match self {
Self::Min => "min",
Self::Max => "max",
}
}
}

impl std::fmt::Display for MinMax {
fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
write!(fmt, "{}", self.as_str())
}
}
2 changes: 2 additions & 0 deletions crates/ruff_linter/src/rules/pylint/rules/mod.rs
Expand Up @@ -21,6 +21,7 @@ pub(crate) use eq_without_hash::*;
pub(crate) use global_at_module_level::*;
pub(crate) use global_statement::*;
pub(crate) use global_variable_not_assigned::*;
pub(crate) use if_stmt_min_max::*;
pub(crate) use import_outside_top_level::*;
pub(crate) use import_private_name::*;
pub(crate) use import_self::*;
Expand Down Expand Up @@ -116,6 +117,7 @@ mod eq_without_hash;
mod global_at_module_level;
mod global_statement;
mod global_variable_not_assigned;
mod if_stmt_min_max;
mod import_outside_top_level;
mod import_private_name;
mod import_self;
Expand Down

0 comments on commit 3194f90

Please sign in to comment.