From 90de108bfa296dfd362b26a037ff2a53abbbef94 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 4 Oct 2023 15:39:05 -0400 Subject: [PATCH] Split up `ast_if.rs` into distinct rule files (#7820) These node-centric rule files are too hard to navigate. Better to have a single file per rule as we do elsewhere. --- .../src/rules/flake8_simplify/rules/ast_if.rs | 1026 ----------------- .../flake8_simplify/rules/collapsible_if.rs | 392 +++++++ .../src/rules/flake8_simplify/rules/fix_if.rs | 132 --- .../if_else_block_instead_of_dict_get.rs | 196 ++++ .../if_else_block_instead_of_dict_lookup.rs | 150 +++ .../rules/if_else_block_instead_of_if_exp.rs | 179 +++ .../rules/if_with_same_arms.rs | 93 ++ .../src/rules/flake8_simplify/rules/mod.rs | 15 +- .../flake8_simplify/rules/needless_bool.rs | 180 +++ 9 files changed, 1202 insertions(+), 1161 deletions(-) delete mode 100644 crates/ruff_linter/src/rules/flake8_simplify/rules/ast_if.rs create mode 100644 crates/ruff_linter/src/rules/flake8_simplify/rules/collapsible_if.rs delete mode 100644 crates/ruff_linter/src/rules/flake8_simplify/rules/fix_if.rs create mode 100644 crates/ruff_linter/src/rules/flake8_simplify/rules/if_else_block_instead_of_dict_get.rs create mode 100644 crates/ruff_linter/src/rules/flake8_simplify/rules/if_else_block_instead_of_dict_lookup.rs create mode 100644 crates/ruff_linter/src/rules/flake8_simplify/rules/if_else_block_instead_of_if_exp.rs create mode 100644 crates/ruff_linter/src/rules/flake8_simplify/rules/if_with_same_arms.rs create mode 100644 crates/ruff_linter/src/rules/flake8_simplify/rules/needless_bool.rs diff --git a/crates/ruff_linter/src/rules/flake8_simplify/rules/ast_if.rs b/crates/ruff_linter/src/rules/flake8_simplify/rules/ast_if.rs deleted file mode 100644 index ba4cf095deb27..0000000000000 --- a/crates/ruff_linter/src/rules/flake8_simplify/rules/ast_if.rs +++ /dev/null @@ -1,1026 +0,0 @@ -use log::error; -use rustc_hash::FxHashSet; - -use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation}; -use ruff_macros::{derive_message_formats, violation}; -use ruff_python_ast::comparable::{ComparableConstant, ComparableExpr, ComparableStmt}; -use ruff_python_ast::helpers::{any_over_expr, contains_effect}; -use ruff_python_ast::node::AnyNodeRef; -use ruff_python_ast::stmt_if::{if_elif_branches, IfElifBranch}; -use ruff_python_ast::{ - self as ast, Arguments, CmpOp, Constant, ElifElseClause, Expr, ExprContext, Identifier, Stmt, -}; -use ruff_python_semantic::SemanticModel; -use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer}; -use ruff_source_file::Locator; -use ruff_text_size::{Ranged, TextRange}; - -use crate::checkers::ast::Checker; -use crate::fix::edits::fits; -use crate::registry::AsRule; -use crate::rules::flake8_simplify::rules::fix_if; - -fn compare_expr(expr1: &ComparableExpr, expr2: &ComparableExpr) -> bool { - expr1.eq(expr2) -} - -fn compare_stmt(stmt1: &ComparableStmt, stmt2: &ComparableStmt) -> bool { - stmt1.eq(stmt2) -} - -/// ## What it does -/// Checks for nested `if` statements that can be collapsed into a single `if` -/// statement. -/// -/// ## Why is this bad? -/// Nesting `if` statements leads to deeper indentation and makes code harder to -/// read. Instead, combine the conditions into a single `if` statement with an -/// `and` operator. -/// -/// ## Example -/// ```python -/// if foo: -/// if bar: -/// ... -/// ``` -/// -/// Use instead: -/// ```python -/// if foo and bar: -/// ... -/// ``` -/// -/// ## References -/// - [Python documentation: The `if` statement](https://docs.python.org/3/reference/compound_stmts.html#the-if-statement) -/// - [Python documentation: Boolean operations](https://docs.python.org/3/reference/expressions.html#boolean-operations) -#[violation] -pub struct CollapsibleIf; - -impl Violation for CollapsibleIf { - const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes; - - #[derive_message_formats] - fn message(&self) -> String { - format!("Use a single `if` statement instead of nested `if` statements") - } - - fn fix_title(&self) -> Option { - Some("Combine `if` statements using `and`".to_string()) - } -} - -/// ## What it does -/// Checks for `if` statements that can be replaced with `bool`. -/// -/// ## Why is this bad? -/// `if` statements that return `True` for a truthy condition and `False` for -/// a falsey condition can be replaced with boolean casts. -/// -/// ## Example -/// ```python -/// if foo: -/// return True -/// else: -/// return False -/// ``` -/// -/// Use instead: -/// ```python -/// return bool(foo) -/// ``` -/// -/// ## References -/// - [Python documentation: Truth Value Testing](https://docs.python.org/3/library/stdtypes.html#truth-value-testing) -#[violation] -pub struct NeedlessBool { - condition: String, -} - -impl Violation for NeedlessBool { - const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes; - - #[derive_message_formats] - fn message(&self) -> String { - let NeedlessBool { condition } = self; - format!("Return the condition `{condition}` directly") - } - - fn fix_title(&self) -> Option { - let NeedlessBool { condition } = self; - Some(format!("Replace with `return {condition}`")) - } -} - -/// ## What it does -/// Checks for three or more consecutive if-statements with direct returns -/// -/// ## Why is this bad? -/// These can be simplified by using a dictionary -/// -/// ## Example -/// ```python -/// if x == 1: -/// return "Hello" -/// elif x == 2: -/// return "Goodbye" -/// else: -/// return "Goodnight" -/// ``` -/// -/// Use instead: -/// ```python -/// return {1: "Hello", 2: "Goodbye"}.get(x, "Goodnight") -/// ``` -#[violation] -pub struct IfElseBlockInsteadOfDictLookup; - -impl Violation for IfElseBlockInsteadOfDictLookup { - #[derive_message_formats] - fn message(&self) -> String { - format!("Use a dictionary instead of consecutive `if` statements") - } -} - -/// ## What it does -/// Check for `if`-`else`-blocks that can be replaced with a ternary operator. -/// -/// ## Why is this bad? -/// `if`-`else`-blocks that assign a value to a variable in both branches can -/// be expressed more concisely by using a ternary operator. -/// -/// ## Example -/// ```python -/// if foo: -/// bar = x -/// else: -/// bar = y -/// ``` -/// -/// Use instead: -/// ```python -/// bar = x if foo else y -/// ``` -/// -/// ## References -/// - [Python documentation: Conditional expressions](https://docs.python.org/3/reference/expressions.html#conditional-expressions) -#[violation] -pub struct IfElseBlockInsteadOfIfExp { - contents: String, -} - -impl Violation for IfElseBlockInsteadOfIfExp { - const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes; - - #[derive_message_formats] - fn message(&self) -> String { - let IfElseBlockInsteadOfIfExp { contents } = self; - format!("Use ternary operator `{contents}` instead of `if`-`else`-block") - } - - fn fix_title(&self) -> Option { - let IfElseBlockInsteadOfIfExp { contents } = self; - Some(format!("Replace `if`-`else`-block with `{contents}`")) - } -} - -/// ## What it does -/// Checks for `if` branches with identical arm bodies. -/// -/// ## Why is this bad? -/// If multiple arms of an `if` statement have the same body, using `or` -/// better signals the intent of the statement. -/// -/// ## Example -/// ```python -/// if x == 1: -/// print("Hello") -/// elif x == 2: -/// print("Hello") -/// ``` -/// -/// Use instead: -/// ```python -/// if x == 1 or x == 2: -/// print("Hello") -/// ``` -#[violation] -pub struct IfWithSameArms; - -impl Violation for IfWithSameArms { - #[derive_message_formats] - fn message(&self) -> String { - format!("Combine `if` branches using logical `or` operator") - } -} - -/// ## What it does -/// Checks for `if` statements that can be replaced with `dict.get` calls. -/// -/// ## Why is this bad? -/// `dict.get()` calls can be used to replace `if` statements that assign a -/// value to a variable in both branches, falling back to a default value if -/// the key is not found. When possible, using `dict.get` is more concise and -/// more idiomatic. -/// -/// ## Example -/// ```python -/// if "bar" in foo: -/// value = foo["bar"] -/// else: -/// value = 0 -/// ``` -/// -/// Use instead: -/// ```python -/// value = foo.get("bar", 0) -/// ``` -/// -/// ## References -/// - [Python documentation: Mapping Types](https://docs.python.org/3/library/stdtypes.html#mapping-types-dict) -#[violation] -pub struct IfElseBlockInsteadOfDictGet { - contents: String, -} - -impl Violation for IfElseBlockInsteadOfDictGet { - const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes; - - #[derive_message_formats] - fn message(&self) -> String { - let IfElseBlockInsteadOfDictGet { contents } = self; - format!("Use `{contents}` instead of an `if` block") - } - - fn fix_title(&self) -> Option { - let IfElseBlockInsteadOfDictGet { contents } = self; - Some(format!("Replace with `{contents}`")) - } -} - -fn is_main_check(expr: &Expr) -> bool { - if let Expr::Compare(ast::ExprCompare { - left, comparators, .. - }) = expr - { - if let Expr::Name(ast::ExprName { id, .. }) = left.as_ref() { - if id == "__name__" { - if let [Expr::Constant(ast::ExprConstant { - value: Constant::Str(ast::StringConstant { value, .. }), - .. - })] = comparators.as_slice() - { - if value == "__main__" { - return true; - } - } - } - } - } - false -} - -/// Find the last nested if statement and return the test expression and the -/// last statement. -/// -/// ```python -/// if xxx: -/// if yyy: -/// # ^^^ returns this expression -/// z = 1 -/// ... -/// ``` -fn find_last_nested_if(body: &[Stmt]) -> Option<&Expr> { - let [Stmt::If(ast::StmtIf { - test, - body: inner_body, - elif_else_clauses, - .. - })] = body - else { - return None; - }; - if !elif_else_clauses.is_empty() { - return None; - } - find_last_nested_if(inner_body).or(Some(test)) -} - -#[derive(Debug, Clone, Copy)] -pub(super) enum NestedIf<'a> { - If(&'a ast::StmtIf), - Elif(&'a ElifElseClause), -} - -impl<'a> NestedIf<'a> { - pub(super) fn body(self) -> &'a [Stmt] { - match self { - NestedIf::If(stmt_if) => &stmt_if.body, - NestedIf::Elif(clause) => &clause.body, - } - } - - pub(super) fn is_elif(self) -> bool { - matches!(self, NestedIf::Elif(..)) - } -} - -impl Ranged for NestedIf<'_> { - fn range(&self) -> TextRange { - match self { - NestedIf::If(stmt_if) => stmt_if.range(), - NestedIf::Elif(clause) => clause.range(), - } - } -} - -impl<'a> From<&NestedIf<'a>> for AnyNodeRef<'a> { - fn from(value: &NestedIf<'a>) -> Self { - match value { - NestedIf::If(stmt_if) => (*stmt_if).into(), - NestedIf::Elif(clause) => (*clause).into(), - } - } -} - -/// Returns the body, the range of the `if` or `elif` and whether the range is for an `if` or `elif` -fn nested_if_body(stmt_if: &ast::StmtIf) -> Option { - let ast::StmtIf { - test, - body, - elif_else_clauses, - .. - } = stmt_if; - - // It must be the last condition, otherwise there could be another `elif` or `else` that only - // depends on the outer of the two conditions - let (test, nested_if) = if let Some(clause) = elif_else_clauses.last() { - if let Some(test) = &clause.test { - (test, NestedIf::Elif(clause)) - } else { - // The last condition is an `else` (different rule) - return None; - } - } else { - (test.as_ref(), NestedIf::If(stmt_if)) - }; - - // The nested if must be the only child, otherwise there is at least one more statement that - // only depends on the outer condition - if body.len() > 1 { - return None; - } - - // Allow `if __name__ == "__main__":` statements. - if is_main_check(test) { - return None; - } - - // Allow `if True:` and `if False:` statements. - if matches!( - test, - Expr::Constant(ast::ExprConstant { - value: Constant::Bool(..), - .. - }) - ) { - return None; - } - - Some(nested_if) -} - -/// SIM102 -pub(crate) fn nested_if_statements( - checker: &mut Checker, - stmt_if: &ast::StmtIf, - parent: Option<&Stmt>, -) { - let Some(nested_if) = nested_if_body(stmt_if) else { - return; - }; - - // Find the deepest nested if-statement, to inform the range. - let Some(test) = find_last_nested_if(nested_if.body()) else { - return; - }; - - // Check if the parent is already emitting a larger diagnostic including this if statement - if let Some(Stmt::If(stmt_if)) = parent { - if let Some(nested_if) = nested_if_body(stmt_if) { - // In addition to repeating the `nested_if_body` and `find_last_nested_if` check, we - // also need to be the first child in the parent - let body = nested_if.body(); - if matches!(&body[0], Stmt::If(inner) if *inner == *stmt_if) - && find_last_nested_if(body).is_some() - { - return; - } - } - } - - let Some(colon) = SimpleTokenizer::starts_at(test.end(), checker.locator().contents()) - .skip_trivia() - .find(|token| token.kind == SimpleTokenKind::Colon) - else { - return; - }; - - let mut diagnostic = Diagnostic::new( - CollapsibleIf, - TextRange::new(nested_if.start(), colon.end()), - ); - if checker.patch(diagnostic.kind.rule()) { - // The fixer preserves comments in the nested body, but removes comments between - // the outer and inner if statements. - if !checker - .indexer() - .comment_ranges() - .intersects(TextRange::new( - nested_if.start(), - nested_if.body()[0].start(), - )) - { - match fix_if::fix_nested_if_statements(checker.locator(), checker.stylist(), nested_if) - { - Ok(edit) => { - if edit.content().map_or(true, |content| { - fits( - content, - (&nested_if).into(), - checker.locator(), - checker.settings.line_length, - checker.settings.tab_size, - ) - }) { - diagnostic.set_fix(Fix::suggested(edit)); - } - } - Err(err) => error!("Failed to fix nested if: {err}"), - } - } - } - checker.diagnostics.push(diagnostic); -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum Bool { - True, - False, -} - -impl From for Bool { - fn from(value: bool) -> Self { - if value { - Bool::True - } else { - Bool::False - } - } -} - -fn is_one_line_return_bool(stmts: &[Stmt]) -> Option { - let [stmt] = stmts else { - return None; - }; - let Stmt::Return(ast::StmtReturn { value, range: _ }) = stmt else { - return None; - }; - let Some(Expr::Constant(ast::ExprConstant { value, .. })) = value.as_deref() else { - return None; - }; - let Constant::Bool(value) = value else { - return None; - }; - Some((*value).into()) -} - -/// SIM103 -pub(crate) fn needless_bool(checker: &mut Checker, stmt: &Stmt) { - let Stmt::If(ast::StmtIf { - test: if_test, - body: if_body, - elif_else_clauses, - range: _, - }) = stmt - else { - return; - }; - // Extract an `if` or `elif` (that returns) followed by an else (that returns the same value) - let (if_test, if_body, else_body, range) = match elif_else_clauses.as_slice() { - // if-else case - [ElifElseClause { - body: else_body, - test: None, - .. - }] => (if_test.as_ref(), if_body, else_body, stmt.range()), - // elif-else case - [.., ElifElseClause { - body: elif_body, - test: Some(elif_test), - range: elif_range, - }, ElifElseClause { - body: else_body, - test: None, - range: else_range, - }] => ( - elif_test, - elif_body, - else_body, - TextRange::new(elif_range.start(), else_range.end()), - ), - _ => return, - }; - - let (Some(if_return), Some(else_return)) = ( - is_one_line_return_bool(if_body), - is_one_line_return_bool(else_body), - ) else { - return; - }; - - // If the branches have the same condition, abort (although the code could be - // simplified). - if if_return == else_return { - return; - } - - let condition = checker.generator().expr(if_test); - let mut diagnostic = Diagnostic::new(NeedlessBool { condition }, range); - if checker.patch(diagnostic.kind.rule()) { - if matches!(if_return, Bool::True) - && matches!(else_return, Bool::False) - && !checker.indexer().has_comments(&range, checker.locator()) - && (if_test.is_compare_expr() || checker.semantic().is_builtin("bool")) - { - if if_test.is_compare_expr() { - // If the condition is a comparison, we can replace it with the condition. - let node = ast::StmtReturn { - value: Some(Box::new(if_test.clone())), - range: TextRange::default(), - }; - diagnostic.set_fix(Fix::suggested(Edit::range_replacement( - checker.generator().stmt(&node.into()), - range, - ))); - } else { - // Otherwise, we need to wrap the condition in a call to `bool`. (We've already - // verified, above, that `bool` is a builtin.) - let node = ast::ExprName { - id: "bool".into(), - ctx: ExprContext::Load, - range: TextRange::default(), - }; - let node1 = ast::ExprCall { - func: Box::new(node.into()), - arguments: Arguments { - args: vec![if_test.clone()], - keywords: vec![], - range: TextRange::default(), - }, - range: TextRange::default(), - }; - let node2 = ast::StmtReturn { - value: Some(Box::new(node1.into())), - range: TextRange::default(), - }; - diagnostic.set_fix(Fix::suggested(Edit::range_replacement( - checker.generator().stmt(&node2.into()), - range, - ))); - }; - } - } - checker.diagnostics.push(diagnostic); -} - -fn ternary(target_var: &Expr, body_value: &Expr, test: &Expr, orelse_value: &Expr) -> Stmt { - let node = ast::ExprIfExp { - test: Box::new(test.clone()), - body: Box::new(body_value.clone()), - orelse: Box::new(orelse_value.clone()), - range: TextRange::default(), - }; - let node1 = ast::StmtAssign { - targets: vec![target_var.clone()], - value: Box::new(node.into()), - range: TextRange::default(), - }; - node1.into() -} - -/// Return `true` if the `Expr` contains a reference to any of the given `${module}.${target}`. -fn contains_call_path(expr: &Expr, targets: &[&[&str]], semantic: &SemanticModel) -> bool { - any_over_expr(expr, &|expr| { - semantic - .resolve_call_path(expr) - .is_some_and(|call_path| targets.iter().any(|target| &call_path.as_slice() == target)) - }) -} - -/// SIM108 -pub(crate) fn use_ternary_operator(checker: &mut Checker, stmt: &Stmt) { - let Stmt::If(ast::StmtIf { - test, - body, - elif_else_clauses, - range: _, - }) = stmt - else { - return; - }; - // `test: None` to only match an `else` clause - let [ElifElseClause { - body: else_body, - test: None, - .. - }] = elif_else_clauses.as_slice() - else { - return; - }; - let [Stmt::Assign(ast::StmtAssign { - targets: body_targets, - value: body_value, - .. - })] = body.as_slice() - else { - return; - }; - let [Stmt::Assign(ast::StmtAssign { - targets: else_targets, - value: else_value, - .. - })] = else_body.as_slice() - else { - return; - }; - let ([body_target], [else_target]) = (body_targets.as_slice(), else_targets.as_slice()) else { - return; - }; - let Expr::Name(ast::ExprName { id: body_id, .. }) = body_target else { - return; - }; - let Expr::Name(ast::ExprName { id: else_id, .. }) = else_target else { - return; - }; - if body_id != else_id { - return; - } - - // Avoid suggesting ternary for `if sys.version_info >= ...`-style and - // `if sys.platform.startswith("...")`-style checks. - let ignored_call_paths: &[&[&str]] = &[&["sys", "version_info"], &["sys", "platform"]]; - if contains_call_path(test, ignored_call_paths, checker.semantic()) { - return; - } - - // Avoid suggesting ternary for `if (yield ...)`-style checks. - // TODO(charlie): Fix precedence handling for yields in generator. - if matches!( - body_value.as_ref(), - Expr::Yield(_) | Expr::YieldFrom(_) | Expr::Await(_) - ) { - return; - } - if matches!( - else_value.as_ref(), - Expr::Yield(_) | Expr::YieldFrom(_) | Expr::Await(_) - ) { - return; - } - - let target_var = &body_target; - let ternary = ternary(target_var, body_value, test, else_value); - let contents = checker.generator().stmt(&ternary); - - // Don't flag if the resulting expression would exceed the maximum line length. - if !fits( - &contents, - stmt.into(), - checker.locator(), - checker.settings.line_length, - checker.settings.tab_size, - ) { - return; - } - - let mut diagnostic = Diagnostic::new( - IfElseBlockInsteadOfIfExp { - contents: contents.clone(), - }, - stmt.range(), - ); - if checker.patch(diagnostic.kind.rule()) { - if !checker.indexer().has_comments(stmt, checker.locator()) { - diagnostic.set_fix(Fix::suggested(Edit::range_replacement( - contents, - stmt.range(), - ))); - } - } - checker.diagnostics.push(diagnostic); -} - -/// Return the [`TextRange`] of an [`IfElifBranch`]'s body (from the end of the test to the end of -/// the body). -fn body_range(branch: &IfElifBranch, locator: &Locator) -> TextRange { - TextRange::new( - locator.line_end(branch.test.end()), - locator.line_end(branch.end()), - ) -} - -/// SIM114 -pub(crate) fn if_with_same_arms(checker: &mut Checker, locator: &Locator, stmt_if: &ast::StmtIf) { - let mut branches_iter = if_elif_branches(stmt_if).peekable(); - while let Some(current_branch) = branches_iter.next() { - let Some(following_branch) = branches_iter.peek() else { - continue; - }; - - // The bodies must have the same code ... - if current_branch.body.len() != following_branch.body.len() { - continue; - } - if !current_branch - .body - .iter() - .zip(following_branch.body.iter()) - .all(|(stmt1, stmt2)| compare_stmt(&stmt1.into(), &stmt2.into())) - { - continue; - } - - // ...and the same comments - let first_comments = checker - .indexer() - .comment_ranges() - .comments_in_range(body_range(¤t_branch, locator)) - .iter() - .map(|range| locator.slice(*range)); - let second_comments = checker - .indexer() - .comment_ranges() - .comments_in_range(body_range(following_branch, locator)) - .iter() - .map(|range| locator.slice(*range)); - if !first_comments.eq(second_comments) { - continue; - } - - checker.diagnostics.push(Diagnostic::new( - IfWithSameArms, - TextRange::new(current_branch.start(), following_branch.end()), - )); - } -} - -/// SIM116 -pub(crate) fn manual_dict_lookup(checker: &mut Checker, stmt_if: &ast::StmtIf) { - // Throughout this rule: - // * Each if or elif statement's test must consist of a constant equality check with the same variable. - // * Each if or elif statement's body must consist of a single `return`. - // * The else clause must be empty, or a single `return`. - let ast::StmtIf { - body, - test, - elif_else_clauses, - .. - } = stmt_if; - - let Expr::Compare(ast::ExprCompare { - left, - ops, - comparators, - range: _, - }) = test.as_ref() - else { - return; - }; - let Expr::Name(ast::ExprName { id: target, .. }) = left.as_ref() else { - return; - }; - if ops != &[CmpOp::Eq] { - return; - } - let [Expr::Constant(ast::ExprConstant { - value: constant, .. - })] = comparators.as_slice() - else { - return; - }; - let [Stmt::Return(ast::StmtReturn { value, range: _ })] = body.as_slice() else { - return; - }; - if value - .as_ref() - .is_some_and(|value| contains_effect(value, |id| checker.semantic().is_builtin(id))) - { - return; - } - - let mut constants: FxHashSet = FxHashSet::default(); - constants.insert(constant.into()); - - for clause in elif_else_clauses { - let ElifElseClause { test, body, .. } = clause; - let [Stmt::Return(ast::StmtReturn { value, range: _ })] = body.as_slice() else { - return; - }; - - match test.as_ref() { - // `else` - None => { - // The else must also be a single effect-free return statement - let [Stmt::Return(ast::StmtReturn { value, range: _ })] = body.as_slice() else { - return; - }; - if value.as_ref().is_some_and(|value| { - contains_effect(value, |id| checker.semantic().is_builtin(id)) - }) { - return; - }; - } - // `elif` - Some(Expr::Compare(ast::ExprCompare { - left, - ops, - comparators, - range: _, - })) => { - let Expr::Name(ast::ExprName { id, .. }) = left.as_ref() else { - return; - }; - if id != target || ops != &[CmpOp::Eq] { - return; - } - let [Expr::Constant(ast::ExprConstant { - value: constant, .. - })] = comparators.as_slice() - else { - return; - }; - - if value.as_ref().is_some_and(|value| { - contains_effect(value, |id| checker.semantic().is_builtin(id)) - }) { - return; - }; - - constants.insert(constant.into()); - } - // Different `elif` - _ => { - return; - } - } - } - - if constants.len() < 3 { - return; - } - - checker.diagnostics.push(Diagnostic::new( - IfElseBlockInsteadOfDictLookup, - stmt_if.range(), - )); -} - -/// SIM401 -pub(crate) fn use_dict_get_with_default(checker: &mut Checker, stmt_if: &ast::StmtIf) { - let ast::StmtIf { - test, - body, - elif_else_clauses, - .. - } = stmt_if; - - let [body_stmt] = body.as_slice() else { - return; - }; - let [ElifElseClause { - body: else_body, - test: None, - .. - }] = elif_else_clauses.as_slice() - else { - return; - }; - let [else_body_stmt] = else_body.as_slice() else { - return; - }; - let Stmt::Assign(ast::StmtAssign { - targets: body_var, - value: body_value, - .. - }) = &body_stmt - else { - return; - }; - let [body_var] = body_var.as_slice() else { - return; - }; - let Stmt::Assign(ast::StmtAssign { - targets: orelse_var, - value: orelse_value, - .. - }) = &else_body_stmt - else { - return; - }; - let [orelse_var] = orelse_var.as_slice() else { - return; - }; - let Expr::Compare(ast::ExprCompare { - left: test_key, - ops, - comparators: test_dict, - range: _, - }) = test.as_ref() - else { - return; - }; - let [test_dict] = test_dict.as_slice() else { - return; - }; - let (expected_var, expected_value, default_var, default_value) = match ops[..] { - [CmpOp::In] => (body_var, body_value, orelse_var, orelse_value.as_ref()), - [CmpOp::NotIn] => (orelse_var, orelse_value, body_var, body_value.as_ref()), - _ => { - return; - } - }; - let Expr::Subscript(ast::ExprSubscript { - value: expected_subscript, - slice: expected_slice, - .. - }) = expected_value.as_ref() - else { - return; - }; - - // Check that the dictionary key, target variables, and dictionary name are all - // equivalent. - if !compare_expr(&expected_slice.into(), &test_key.into()) - || !compare_expr(&expected_var.into(), &default_var.into()) - || !compare_expr(&test_dict.into(), &expected_subscript.into()) - { - return; - } - - // Check that the default value is not "complex". - if contains_effect(default_value, |id| checker.semantic().is_builtin(id)) { - return; - } - - let node = default_value.clone(); - let node1 = *test_key.clone(); - let node2 = ast::ExprAttribute { - value: expected_subscript.clone(), - attr: Identifier::new("get".to_string(), TextRange::default()), - ctx: ExprContext::Load, - range: TextRange::default(), - }; - let node3 = ast::ExprCall { - func: Box::new(node2.into()), - arguments: Arguments { - args: vec![node1, node], - keywords: vec![], - range: TextRange::default(), - }, - range: TextRange::default(), - }; - let node4 = expected_var.clone(); - let node5 = ast::StmtAssign { - targets: vec![node4], - value: Box::new(node3.into()), - range: TextRange::default(), - }; - let contents = checker.generator().stmt(&node5.into()); - - // Don't flag if the resulting expression would exceed the maximum line length. - if !fits( - &contents, - stmt_if.into(), - checker.locator(), - checker.settings.line_length, - checker.settings.tab_size, - ) { - return; - } - - let mut diagnostic = Diagnostic::new( - IfElseBlockInsteadOfDictGet { - contents: contents.clone(), - }, - stmt_if.range(), - ); - if checker.patch(diagnostic.kind.rule()) { - if !checker.indexer().has_comments(stmt_if, checker.locator()) { - diagnostic.set_fix(Fix::suggested(Edit::range_replacement( - contents, - stmt_if.range(), - ))); - } - } - checker.diagnostics.push(diagnostic); -} diff --git a/crates/ruff_linter/src/rules/flake8_simplify/rules/collapsible_if.rs b/crates/ruff_linter/src/rules/flake8_simplify/rules/collapsible_if.rs new file mode 100644 index 0000000000000..e8601624d0744 --- /dev/null +++ b/crates/ruff_linter/src/rules/flake8_simplify/rules/collapsible_if.rs @@ -0,0 +1,392 @@ +use std::borrow::Cow; + +use anyhow::{bail, Result}; +use libcst_native::ParenthesizedNode; +use log::error; + +use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation}; +use ruff_macros::{derive_message_formats, violation}; +use ruff_python_ast::node::AnyNodeRef; +use ruff_python_ast::{self as ast, whitespace, Constant, ElifElseClause, Expr, Stmt}; +use ruff_python_codegen::Stylist; +use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer}; +use ruff_source_file::Locator; +use ruff_text_size::{Ranged, TextRange}; + +use crate::checkers::ast::Checker; +use crate::cst::helpers::space; +use crate::cst::matchers::{match_function_def, match_if, match_indented_block, match_statement}; +use crate::fix::codemods::CodegenStylist; +use crate::fix::edits::fits; +use crate::registry::AsRule; + +/// ## What it does +/// Checks for nested `if` statements that can be collapsed into a single `if` +/// statement. +/// +/// ## Why is this bad? +/// Nesting `if` statements leads to deeper indentation and makes code harder to +/// read. Instead, combine the conditions into a single `if` statement with an +/// `and` operator. +/// +/// ## Example +/// ```python +/// if foo: +/// if bar: +/// ... +/// ``` +/// +/// Use instead: +/// ```python +/// if foo and bar: +/// ... +/// ``` +/// +/// ## References +/// - [Python documentation: The `if` statement](https://docs.python.org/3/reference/compound_stmts.html#the-if-statement) +/// - [Python documentation: Boolean operations](https://docs.python.org/3/reference/expressions.html#boolean-operations) +#[violation] +pub struct CollapsibleIf; + +impl Violation for CollapsibleIf { + const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes; + + #[derive_message_formats] + fn message(&self) -> String { + format!("Use a single `if` statement instead of nested `if` statements") + } + + fn fix_title(&self) -> Option { + Some("Combine `if` statements using `and`".to_string()) + } +} + +/// SIM102 +pub(crate) fn nested_if_statements( + checker: &mut Checker, + stmt_if: &ast::StmtIf, + parent: Option<&Stmt>, +) { + let Some(nested_if) = nested_if_body(stmt_if) else { + return; + }; + + // Find the deepest nested if-statement, to inform the range. + let Some(test) = find_last_nested_if(nested_if.body()) else { + return; + }; + + // Check if the parent is already emitting a larger diagnostic including this if statement + if let Some(Stmt::If(stmt_if)) = parent { + if let Some(nested_if) = nested_if_body(stmt_if) { + // In addition to repeating the `nested_if_body` and `find_last_nested_if` check, we + // also need to be the first child in the parent + let body = nested_if.body(); + if matches!(&body[0], Stmt::If(inner) if *inner == *stmt_if) + && find_last_nested_if(body).is_some() + { + return; + } + } + } + + let Some(colon) = SimpleTokenizer::starts_at(test.end(), checker.locator().contents()) + .skip_trivia() + .find(|token| token.kind == SimpleTokenKind::Colon) + else { + return; + }; + + let mut diagnostic = Diagnostic::new( + CollapsibleIf, + TextRange::new(nested_if.start(), colon.end()), + ); + if checker.patch(diagnostic.kind.rule()) { + // The fixer preserves comments in the nested body, but removes comments between + // the outer and inner if statements. + if !checker + .indexer() + .comment_ranges() + .intersects(TextRange::new( + nested_if.start(), + nested_if.body()[0].start(), + )) + { + match collapse_nested_if(checker.locator(), checker.stylist(), nested_if) { + Ok(edit) => { + if edit.content().map_or(true, |content| { + fits( + content, + (&nested_if).into(), + checker.locator(), + checker.settings.line_length, + checker.settings.tab_size, + ) + }) { + diagnostic.set_fix(Fix::suggested(edit)); + } + } + Err(err) => error!("Failed to fix nested if: {err}"), + } + } + } + checker.diagnostics.push(diagnostic); +} + +#[derive(Debug, Clone, Copy)] +pub(super) enum NestedIf<'a> { + If(&'a ast::StmtIf), + Elif(&'a ElifElseClause), +} + +impl<'a> NestedIf<'a> { + pub(super) fn body(self) -> &'a [Stmt] { + match self { + NestedIf::If(stmt_if) => &stmt_if.body, + NestedIf::Elif(clause) => &clause.body, + } + } + + pub(super) fn is_elif(self) -> bool { + matches!(self, NestedIf::Elif(..)) + } +} + +impl Ranged for NestedIf<'_> { + fn range(&self) -> TextRange { + match self { + NestedIf::If(stmt_if) => stmt_if.range(), + NestedIf::Elif(clause) => clause.range(), + } + } +} + +impl<'a> From<&NestedIf<'a>> for AnyNodeRef<'a> { + fn from(value: &NestedIf<'a>) -> Self { + match value { + NestedIf::If(stmt_if) => (*stmt_if).into(), + NestedIf::Elif(clause) => (*clause).into(), + } + } +} + +/// Returns the body, the range of the `if` or `elif` and whether the range is for an `if` or `elif` +fn nested_if_body(stmt_if: &ast::StmtIf) -> Option { + let ast::StmtIf { + test, + body, + elif_else_clauses, + .. + } = stmt_if; + + // It must be the last condition, otherwise there could be another `elif` or `else` that only + // depends on the outer of the two conditions + let (test, nested_if) = if let Some(clause) = elif_else_clauses.last() { + if let Some(test) = &clause.test { + (test, NestedIf::Elif(clause)) + } else { + // The last condition is an `else` (different rule) + return None; + } + } else { + (test.as_ref(), NestedIf::If(stmt_if)) + }; + + // The nested if must be the only child, otherwise there is at least one more statement that + // only depends on the outer condition + if body.len() > 1 { + return None; + } + + // Allow `if __name__ == "__main__":` statements. + if is_main_check(test) { + return None; + } + + // Allow `if True:` and `if False:` statements. + if matches!( + test, + Expr::Constant(ast::ExprConstant { + value: Constant::Bool(..), + .. + }) + ) { + return None; + } + + Some(nested_if) +} + +/// Find the last nested if statement and return the test expression and the +/// last statement. +/// +/// ```python +/// if xxx: +/// if yyy: +/// # ^^^ returns this expression +/// z = 1 +/// ... +/// ``` +fn find_last_nested_if(body: &[Stmt]) -> Option<&Expr> { + let [Stmt::If(ast::StmtIf { + test, + body: inner_body, + elif_else_clauses, + .. + })] = body + else { + return None; + }; + if !elif_else_clauses.is_empty() { + return None; + } + find_last_nested_if(inner_body).or(Some(test)) +} + +/// Returns `true` if an expression is an `if __name__ == "__main__":` check. +fn is_main_check(expr: &Expr) -> bool { + if let Expr::Compare(ast::ExprCompare { + left, comparators, .. + }) = expr + { + if let Expr::Name(ast::ExprName { id, .. }) = left.as_ref() { + if id == "__name__" { + if let [Expr::Constant(ast::ExprConstant { + value: Constant::Str(ast::StringConstant { value, .. }), + .. + })] = comparators.as_slice() + { + if value == "__main__" { + return true; + } + } + } + } + } + false +} + +fn parenthesize_and_operand(expr: libcst_native::Expression) -> libcst_native::Expression { + match &expr { + _ if !expr.lpar().is_empty() => expr, + libcst_native::Expression::BooleanOperation(boolean_operation) + if matches!( + boolean_operation.operator, + libcst_native::BooleanOp::Or { .. } + ) => + { + expr.with_parens( + libcst_native::LeftParen::default(), + libcst_native::RightParen::default(), + ) + } + libcst_native::Expression::IfExp(_) + | libcst_native::Expression::Lambda(_) + | libcst_native::Expression::NamedExpr(_) => expr.with_parens( + libcst_native::LeftParen::default(), + libcst_native::RightParen::default(), + ), + _ => expr, + } +} + +/// Convert `if a: if b:` to `if a and b:`. +pub(super) fn collapse_nested_if( + locator: &Locator, + stylist: &Stylist, + nested_if: NestedIf, +) -> Result { + // Infer the indentation of the outer block. + let Some(outer_indent) = whitespace::indentation(locator, &nested_if) else { + bail!("Unable to fix multiline statement"); + }; + + // Extract the module text. + let contents = locator.lines(nested_if.range()); + + // If this is an `elif`, we have to remove the `elif` keyword for now. (We'll + // restore the `el` later on.) + let module_text = if nested_if.is_elif() { + Cow::Owned(contents.replacen("elif", "if", 1)) + } else { + Cow::Borrowed(contents) + }; + + // If the block is indented, "embed" it in a function definition, to preserve + // indentation while retaining valid source code. (We'll strip the prefix later + // on.) + let module_text = if outer_indent.is_empty() { + module_text + } else { + Cow::Owned(format!( + "def f():{}{module_text}", + stylist.line_ending().as_str() + )) + }; + + // Parse the CST. + let mut tree = match_statement(&module_text)?; + + let statement = if outer_indent.is_empty() { + &mut tree + } else { + let embedding = match_function_def(&mut tree)?; + + let indented_block = match_indented_block(&mut embedding.body)?; + indented_block.indent = Some(outer_indent); + + let Some(statement) = indented_block.body.first_mut() else { + bail!("Expected indented block to have at least one statement") + }; + statement + }; + + let outer_if = match_if(statement)?; + + let libcst_native::If { + body: libcst_native::Suite::IndentedBlock(ref mut outer_body), + orelse: None, + .. + } = outer_if + else { + bail!("Expected outer if to have indented body and no else") + }; + + let [libcst_native::Statement::Compound(libcst_native::CompoundStatement::If( + inner_if @ libcst_native::If { orelse: None, .. }, + ))] = &mut *outer_body.body + else { + bail!("Expected one inner if statement"); + }; + + outer_if.test = + libcst_native::Expression::BooleanOperation(Box::new(libcst_native::BooleanOperation { + left: Box::new(parenthesize_and_operand(outer_if.test.clone())), + operator: libcst_native::BooleanOp::And { + whitespace_before: space(), + whitespace_after: space(), + }, + right: Box::new(parenthesize_and_operand(inner_if.test.clone())), + lpar: vec![], + rpar: vec![], + })); + outer_if.body = inner_if.body.clone(); + + // Reconstruct and reformat the code. + let module_text = tree.codegen_stylist(stylist); + let module_text = if outer_indent.is_empty() { + &module_text + } else { + module_text + .strip_prefix(&format!("def f():{}", stylist.line_ending().as_str())) + .unwrap() + }; + let contents = if nested_if.is_elif() { + Cow::Owned(module_text.replacen("if", "elif", 1)) + } else { + Cow::Borrowed(module_text) + }; + + let range = locator.lines_range(nested_if.range()); + Ok(Edit::range_replacement(contents.to_string(), range)) +} diff --git a/crates/ruff_linter/src/rules/flake8_simplify/rules/fix_if.rs b/crates/ruff_linter/src/rules/flake8_simplify/rules/fix_if.rs deleted file mode 100644 index bf775fb365280..0000000000000 --- a/crates/ruff_linter/src/rules/flake8_simplify/rules/fix_if.rs +++ /dev/null @@ -1,132 +0,0 @@ -use std::borrow::Cow; - -use anyhow::{bail, Result}; -use libcst_native::{ - BooleanOp, BooleanOperation, CompoundStatement, Expression, If, LeftParen, ParenthesizedNode, - RightParen, Statement, Suite, -}; - -use ruff_diagnostics::Edit; -use ruff_python_ast::whitespace; -use ruff_python_codegen::Stylist; -use ruff_source_file::Locator; -use ruff_text_size::Ranged; - -use crate::cst::helpers::space; -use crate::cst::matchers::{match_function_def, match_if, match_indented_block, match_statement}; -use crate::fix::codemods::CodegenStylist; -use crate::rules::flake8_simplify::rules::ast_if::NestedIf; - -fn parenthesize_and_operand(expr: Expression) -> Expression { - match &expr { - _ if !expr.lpar().is_empty() => expr, - Expression::BooleanOperation(boolean_operation) - if matches!(boolean_operation.operator, BooleanOp::Or { .. }) => - { - expr.with_parens(LeftParen::default(), RightParen::default()) - } - Expression::IfExp(_) | Expression::Lambda(_) | Expression::NamedExpr(_) => { - expr.with_parens(LeftParen::default(), RightParen::default()) - } - _ => expr, - } -} - -/// (SIM102) Convert `if a: if b:` to `if a and b:`. -pub(super) fn fix_nested_if_statements( - locator: &Locator, - stylist: &Stylist, - nested_if: NestedIf, -) -> Result { - // Infer the indentation of the outer block. - let Some(outer_indent) = whitespace::indentation(locator, &nested_if) else { - bail!("Unable to fix multiline statement"); - }; - - // Extract the module text. - let contents = locator.lines(nested_if.range()); - - // If this is an `elif`, we have to remove the `elif` keyword for now. (We'll - // restore the `el` later on.) - let module_text = if nested_if.is_elif() { - Cow::Owned(contents.replacen("elif", "if", 1)) - } else { - Cow::Borrowed(contents) - }; - - // If the block is indented, "embed" it in a function definition, to preserve - // indentation while retaining valid source code. (We'll strip the prefix later - // on.) - let module_text = if outer_indent.is_empty() { - module_text - } else { - Cow::Owned(format!( - "def f():{}{module_text}", - stylist.line_ending().as_str() - )) - }; - - // Parse the CST. - let mut tree = match_statement(&module_text)?; - - let statement = if outer_indent.is_empty() { - &mut tree - } else { - let embedding = match_function_def(&mut tree)?; - - let indented_block = match_indented_block(&mut embedding.body)?; - indented_block.indent = Some(outer_indent); - - let Some(statement) = indented_block.body.first_mut() else { - bail!("Expected indented block to have at least one statement") - }; - statement - }; - - let outer_if = match_if(statement)?; - - let If { - body: Suite::IndentedBlock(ref mut outer_body), - orelse: None, - .. - } = outer_if - else { - bail!("Expected outer if to have indented body and no else") - }; - - let [Statement::Compound(CompoundStatement::If(inner_if @ If { orelse: None, .. }))] = - &mut *outer_body.body - else { - bail!("Expected one inner if statement"); - }; - - outer_if.test = Expression::BooleanOperation(Box::new(BooleanOperation { - left: Box::new(parenthesize_and_operand(outer_if.test.clone())), - operator: BooleanOp::And { - whitespace_before: space(), - whitespace_after: space(), - }, - right: Box::new(parenthesize_and_operand(inner_if.test.clone())), - lpar: vec![], - rpar: vec![], - })); - outer_if.body = inner_if.body.clone(); - - // Reconstruct and reformat the code. - let module_text = tree.codegen_stylist(stylist); - let module_text = if outer_indent.is_empty() { - &module_text - } else { - module_text - .strip_prefix(&format!("def f():{}", stylist.line_ending().as_str())) - .unwrap() - }; - let contents = if nested_if.is_elif() { - Cow::Owned(module_text.replacen("if", "elif", 1)) - } else { - Cow::Borrowed(module_text) - }; - - let range = locator.lines_range(nested_if.range()); - Ok(Edit::range_replacement(contents.to_string(), range)) -} diff --git a/crates/ruff_linter/src/rules/flake8_simplify/rules/if_else_block_instead_of_dict_get.rs b/crates/ruff_linter/src/rules/flake8_simplify/rules/if_else_block_instead_of_dict_get.rs new file mode 100644 index 0000000000000..d70348801664d --- /dev/null +++ b/crates/ruff_linter/src/rules/flake8_simplify/rules/if_else_block_instead_of_dict_get.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::helpers::contains_effect; +use ruff_python_ast::{ + self as ast, Arguments, CmpOp, ElifElseClause, Expr, ExprContext, Identifier, Stmt, +}; +use ruff_text_size::{Ranged, TextRange}; + +use crate::checkers::ast::Checker; +use crate::fix::edits::fits; +use crate::registry::AsRule; + +/// ## What it does +/// Checks for `if` statements that can be replaced with `dict.get` calls. +/// +/// ## Why is this bad? +/// `dict.get()` calls can be used to replace `if` statements that assign a +/// value to a variable in both branches, falling back to a default value if +/// the key is not found. When possible, using `dict.get` is more concise and +/// more idiomatic. +/// +/// ## Example +/// ```python +/// if "bar" in foo: +/// value = foo["bar"] +/// else: +/// value = 0 +/// ``` +/// +/// Use instead: +/// ```python +/// value = foo.get("bar", 0) +/// ``` +/// +/// ## References +/// - [Python documentation: Mapping Types](https://docs.python.org/3/library/stdtypes.html#mapping-types-dict) +#[violation] +pub struct IfElseBlockInsteadOfDictGet { + contents: String, +} + +impl Violation for IfElseBlockInsteadOfDictGet { + const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes; + + #[derive_message_formats] + fn message(&self) -> String { + let IfElseBlockInsteadOfDictGet { contents } = self; + format!("Use `{contents}` instead of an `if` block") + } + + fn fix_title(&self) -> Option { + let IfElseBlockInsteadOfDictGet { contents } = self; + Some(format!("Replace with `{contents}`")) + } +} + +/// SIM401 +pub(crate) fn use_dict_get_with_default(checker: &mut Checker, stmt_if: &ast::StmtIf) { + let ast::StmtIf { + test, + body, + elif_else_clauses, + .. + } = stmt_if; + + let [body_stmt] = body.as_slice() else { + return; + }; + let [ElifElseClause { + body: else_body, + test: None, + .. + }] = elif_else_clauses.as_slice() + else { + return; + }; + let [else_body_stmt] = else_body.as_slice() else { + return; + }; + let Stmt::Assign(ast::StmtAssign { + targets: body_var, + value: body_value, + .. + }) = &body_stmt + else { + return; + }; + let [body_var] = body_var.as_slice() else { + return; + }; + let Stmt::Assign(ast::StmtAssign { + targets: orelse_var, + value: orelse_value, + .. + }) = &else_body_stmt + else { + return; + }; + let [orelse_var] = orelse_var.as_slice() else { + return; + }; + let Expr::Compare(ast::ExprCompare { + left: test_key, + ops, + comparators: test_dict, + range: _, + }) = test.as_ref() + else { + return; + }; + let [test_dict] = test_dict.as_slice() else { + return; + }; + let (expected_var, expected_value, default_var, default_value) = match ops[..] { + [CmpOp::In] => (body_var, body_value, orelse_var, orelse_value.as_ref()), + [CmpOp::NotIn] => (orelse_var, orelse_value, body_var, body_value.as_ref()), + _ => { + return; + } + }; + let Expr::Subscript(ast::ExprSubscript { + value: expected_subscript, + slice: expected_slice, + .. + }) = expected_value.as_ref() + else { + return; + }; + + // Check that the dictionary key, target variables, and dictionary name are all + // equivalent. + if ComparableExpr::from(expected_slice) != ComparableExpr::from(test_key) + || ComparableExpr::from(expected_var) != ComparableExpr::from(default_var) + || ComparableExpr::from(test_dict) != ComparableExpr::from(expected_subscript) + { + return; + } + + // Check that the default value is not "complex". + if contains_effect(default_value, |id| checker.semantic().is_builtin(id)) { + return; + } + + let node = default_value.clone(); + let node1 = *test_key.clone(); + let node2 = ast::ExprAttribute { + value: expected_subscript.clone(), + attr: Identifier::new("get".to_string(), TextRange::default()), + ctx: ExprContext::Load, + range: TextRange::default(), + }; + let node3 = ast::ExprCall { + func: Box::new(node2.into()), + arguments: Arguments { + args: vec![node1, node], + keywords: vec![], + range: TextRange::default(), + }, + range: TextRange::default(), + }; + let node4 = expected_var.clone(); + let node5 = ast::StmtAssign { + targets: vec![node4], + value: Box::new(node3.into()), + range: TextRange::default(), + }; + let contents = checker.generator().stmt(&node5.into()); + + // Don't flag if the resulting expression would exceed the maximum line length. + if !fits( + &contents, + stmt_if.into(), + checker.locator(), + checker.settings.line_length, + checker.settings.tab_size, + ) { + return; + } + + let mut diagnostic = Diagnostic::new( + IfElseBlockInsteadOfDictGet { + contents: contents.clone(), + }, + stmt_if.range(), + ); + if checker.patch(diagnostic.kind.rule()) { + if !checker.indexer().has_comments(stmt_if, checker.locator()) { + diagnostic.set_fix(Fix::suggested(Edit::range_replacement( + contents, + stmt_if.range(), + ))); + } + } + checker.diagnostics.push(diagnostic); +} diff --git a/crates/ruff_linter/src/rules/flake8_simplify/rules/if_else_block_instead_of_dict_lookup.rs b/crates/ruff_linter/src/rules/flake8_simplify/rules/if_else_block_instead_of_dict_lookup.rs new file mode 100644 index 0000000000000..a7afd23c98583 --- /dev/null +++ b/crates/ruff_linter/src/rules/flake8_simplify/rules/if_else_block_instead_of_dict_lookup.rs @@ -0,0 +1,150 @@ +use rustc_hash::FxHashSet; + +use ruff_diagnostics::{Diagnostic, Violation}; +use ruff_macros::{derive_message_formats, violation}; +use ruff_python_ast::comparable::ComparableConstant; +use ruff_python_ast::helpers::contains_effect; +use ruff_python_ast::{self as ast, CmpOp, ElifElseClause, Expr, Stmt}; +use ruff_text_size::Ranged; + +use crate::checkers::ast::Checker; + +/// ## What it does +/// Checks for three or more consecutive if-statements with direct returns +/// +/// ## Why is this bad? +/// These can be simplified by using a dictionary +/// +/// ## Example +/// ```python +/// if x == 1: +/// return "Hello" +/// elif x == 2: +/// return "Goodbye" +/// else: +/// return "Goodnight" +/// ``` +/// +/// Use instead: +/// ```python +/// return {1: "Hello", 2: "Goodbye"}.get(x, "Goodnight") +/// ``` +#[violation] +pub struct IfElseBlockInsteadOfDictLookup; + +impl Violation for IfElseBlockInsteadOfDictLookup { + #[derive_message_formats] + fn message(&self) -> String { + format!("Use a dictionary instead of consecutive `if` statements") + } +} +/// SIM116 +pub(crate) fn manual_dict_lookup(checker: &mut Checker, stmt_if: &ast::StmtIf) { + // Throughout this rule: + // * Each if or elif statement's test must consist of a constant equality check with the same variable. + // * Each if or elif statement's body must consist of a single `return`. + // * The else clause must be empty, or a single `return`. + let ast::StmtIf { + body, + test, + elif_else_clauses, + .. + } = stmt_if; + + let Expr::Compare(ast::ExprCompare { + left, + ops, + comparators, + range: _, + }) = test.as_ref() + else { + return; + }; + let Expr::Name(ast::ExprName { id: target, .. }) = left.as_ref() else { + return; + }; + if ops != &[CmpOp::Eq] { + return; + } + let [Expr::Constant(ast::ExprConstant { + value: constant, .. + })] = comparators.as_slice() + else { + return; + }; + let [Stmt::Return(ast::StmtReturn { value, range: _ })] = body.as_slice() else { + return; + }; + if value + .as_ref() + .is_some_and(|value| contains_effect(value, |id| checker.semantic().is_builtin(id))) + { + return; + } + + let mut constants: FxHashSet = FxHashSet::default(); + constants.insert(constant.into()); + + for clause in elif_else_clauses { + let ElifElseClause { test, body, .. } = clause; + let [Stmt::Return(ast::StmtReturn { value, range: _ })] = body.as_slice() else { + return; + }; + + match test.as_ref() { + // `else` + None => { + // The else must also be a single effect-free return statement + let [Stmt::Return(ast::StmtReturn { value, range: _ })] = body.as_slice() else { + return; + }; + if value.as_ref().is_some_and(|value| { + contains_effect(value, |id| checker.semantic().is_builtin(id)) + }) { + return; + }; + } + // `elif` + Some(Expr::Compare(ast::ExprCompare { + left, + ops, + comparators, + range: _, + })) => { + let Expr::Name(ast::ExprName { id, .. }) = left.as_ref() else { + return; + }; + if id != target || ops != &[CmpOp::Eq] { + return; + } + let [Expr::Constant(ast::ExprConstant { + value: constant, .. + })] = comparators.as_slice() + else { + return; + }; + + if value.as_ref().is_some_and(|value| { + contains_effect(value, |id| checker.semantic().is_builtin(id)) + }) { + return; + }; + + constants.insert(constant.into()); + } + // Different `elif` + _ => { + return; + } + } + } + + if constants.len() < 3 { + return; + } + + checker.diagnostics.push(Diagnostic::new( + IfElseBlockInsteadOfDictLookup, + stmt_if.range(), + )); +} diff --git a/crates/ruff_linter/src/rules/flake8_simplify/rules/if_else_block_instead_of_if_exp.rs b/crates/ruff_linter/src/rules/flake8_simplify/rules/if_else_block_instead_of_if_exp.rs new file mode 100644 index 0000000000000..fce5712152e08 --- /dev/null +++ b/crates/ruff_linter/src/rules/flake8_simplify/rules/if_else_block_instead_of_if_exp.rs @@ -0,0 +1,179 @@ +use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation}; +use ruff_macros::{derive_message_formats, violation}; +use ruff_python_ast::helpers::any_over_expr; +use ruff_python_ast::{self as ast, ElifElseClause, Expr, Stmt}; +use ruff_python_semantic::SemanticModel; +use ruff_text_size::{Ranged, TextRange}; + +use crate::checkers::ast::Checker; +use crate::fix::edits::fits; +use crate::registry::AsRule; + +/// ## What it does +/// Check for `if`-`else`-blocks that can be replaced with a ternary operator. +/// +/// ## Why is this bad? +/// `if`-`else`-blocks that assign a value to a variable in both branches can +/// be expressed more concisely by using a ternary operator. +/// +/// ## Example +/// ```python +/// if foo: +/// bar = x +/// else: +/// bar = y +/// ``` +/// +/// Use instead: +/// ```python +/// bar = x if foo else y +/// ``` +/// +/// ## References +/// - [Python documentation: Conditional expressions](https://docs.python.org/3/reference/expressions.html#conditional-expressions) +#[violation] +pub struct IfElseBlockInsteadOfIfExp { + contents: String, +} + +impl Violation for IfElseBlockInsteadOfIfExp { + const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes; + + #[derive_message_formats] + fn message(&self) -> String { + let IfElseBlockInsteadOfIfExp { contents } = self; + format!("Use ternary operator `{contents}` instead of `if`-`else`-block") + } + + fn fix_title(&self) -> Option { + let IfElseBlockInsteadOfIfExp { contents } = self; + Some(format!("Replace `if`-`else`-block with `{contents}`")) + } +} + +/// SIM108 +pub(crate) fn use_ternary_operator(checker: &mut Checker, stmt: &Stmt) { + let Stmt::If(ast::StmtIf { + test, + body, + elif_else_clauses, + range: _, + }) = stmt + else { + return; + }; + // `test: None` to only match an `else` clause + let [ElifElseClause { + body: else_body, + test: None, + .. + }] = elif_else_clauses.as_slice() + else { + return; + }; + let [Stmt::Assign(ast::StmtAssign { + targets: body_targets, + value: body_value, + .. + })] = body.as_slice() + else { + return; + }; + let [Stmt::Assign(ast::StmtAssign { + targets: else_targets, + value: else_value, + .. + })] = else_body.as_slice() + else { + return; + }; + let ([body_target], [else_target]) = (body_targets.as_slice(), else_targets.as_slice()) else { + return; + }; + let Expr::Name(ast::ExprName { id: body_id, .. }) = body_target else { + return; + }; + let Expr::Name(ast::ExprName { id: else_id, .. }) = else_target else { + return; + }; + if body_id != else_id { + return; + } + + // Avoid suggesting ternary for `if sys.version_info >= ...`-style and + // `if sys.platform.startswith("...")`-style checks. + let ignored_call_paths: &[&[&str]] = &[&["sys", "version_info"], &["sys", "platform"]]; + if contains_call_path(test, ignored_call_paths, checker.semantic()) { + return; + } + + // Avoid suggesting ternary for `if (yield ...)`-style checks. + // TODO(charlie): Fix precedence handling for yields in generator. + if matches!( + body_value.as_ref(), + Expr::Yield(_) | Expr::YieldFrom(_) | Expr::Await(_) + ) { + return; + } + if matches!( + else_value.as_ref(), + Expr::Yield(_) | Expr::YieldFrom(_) | Expr::Await(_) + ) { + return; + } + + let target_var = &body_target; + let ternary = ternary(target_var, body_value, test, else_value); + let contents = checker.generator().stmt(&ternary); + + // Don't flag if the resulting expression would exceed the maximum line length. + if !fits( + &contents, + stmt.into(), + checker.locator(), + checker.settings.line_length, + checker.settings.tab_size, + ) { + return; + } + + let mut diagnostic = Diagnostic::new( + IfElseBlockInsteadOfIfExp { + contents: contents.clone(), + }, + stmt.range(), + ); + if checker.patch(diagnostic.kind.rule()) { + if !checker.indexer().has_comments(stmt, checker.locator()) { + diagnostic.set_fix(Fix::suggested(Edit::range_replacement( + contents, + stmt.range(), + ))); + } + } + checker.diagnostics.push(diagnostic); +} + +/// Return `true` if the `Expr` contains a reference to any of the given `${module}.${target}`. +fn contains_call_path(expr: &Expr, targets: &[&[&str]], semantic: &SemanticModel) -> bool { + any_over_expr(expr, &|expr| { + semantic + .resolve_call_path(expr) + .is_some_and(|call_path| targets.iter().any(|target| &call_path.as_slice() == target)) + }) +} + +fn ternary(target_var: &Expr, body_value: &Expr, test: &Expr, orelse_value: &Expr) -> Stmt { + let node = ast::ExprIfExp { + test: Box::new(test.clone()), + body: Box::new(body_value.clone()), + orelse: Box::new(orelse_value.clone()), + range: TextRange::default(), + }; + let node1 = ast::StmtAssign { + targets: vec![target_var.clone()], + value: Box::new(node.into()), + range: TextRange::default(), + }; + node1.into() +} diff --git a/crates/ruff_linter/src/rules/flake8_simplify/rules/if_with_same_arms.rs b/crates/ruff_linter/src/rules/flake8_simplify/rules/if_with_same_arms.rs new file mode 100644 index 0000000000000..9ffc34b236a00 --- /dev/null +++ b/crates/ruff_linter/src/rules/flake8_simplify/rules/if_with_same_arms.rs @@ -0,0 +1,93 @@ +use ruff_diagnostics::{Diagnostic, Violation}; +use ruff_macros::{derive_message_formats, violation}; +use ruff_python_ast as ast; +use ruff_python_ast::comparable::ComparableStmt; +use ruff_python_ast::stmt_if::{if_elif_branches, IfElifBranch}; +use ruff_source_file::Locator; +use ruff_text_size::{Ranged, TextRange}; + +use crate::checkers::ast::Checker; + +/// ## What it does +/// Checks for `if` branches with identical arm bodies. +/// +/// ## Why is this bad? +/// If multiple arms of an `if` statement have the same body, using `or` +/// better signals the intent of the statement. +/// +/// ## Example +/// ```python +/// if x == 1: +/// print("Hello") +/// elif x == 2: +/// print("Hello") +/// ``` +/// +/// Use instead: +/// ```python +/// if x == 1 or x == 2: +/// print("Hello") +/// ``` +#[violation] +pub struct IfWithSameArms; + +impl Violation for IfWithSameArms { + #[derive_message_formats] + fn message(&self) -> String { + format!("Combine `if` branches using logical `or` operator") + } +} + +/// SIM114 +pub(crate) fn if_with_same_arms(checker: &mut Checker, locator: &Locator, stmt_if: &ast::StmtIf) { + let mut branches_iter = if_elif_branches(stmt_if).peekable(); + while let Some(current_branch) = branches_iter.next() { + let Some(following_branch) = branches_iter.peek() else { + continue; + }; + + // The bodies must have the same code ... + if current_branch.body.len() != following_branch.body.len() { + continue; + } + if !current_branch + .body + .iter() + .zip(following_branch.body.iter()) + .all(|(stmt1, stmt2)| ComparableStmt::from(stmt1) == ComparableStmt::from(stmt2)) + { + continue; + } + + // ...and the same comments + let first_comments = checker + .indexer() + .comment_ranges() + .comments_in_range(body_range(¤t_branch, locator)) + .iter() + .map(|range| locator.slice(*range)); + let second_comments = checker + .indexer() + .comment_ranges() + .comments_in_range(body_range(following_branch, locator)) + .iter() + .map(|range| locator.slice(*range)); + if !first_comments.eq(second_comments) { + continue; + } + + checker.diagnostics.push(Diagnostic::new( + IfWithSameArms, + TextRange::new(current_branch.start(), following_branch.end()), + )); + } +} + +/// Return the [`TextRange`] of an [`IfElifBranch`]'s body (from the end of the test to the end of +/// the body). +fn body_range(branch: &IfElifBranch, locator: &Locator) -> TextRange { + TextRange::new( + locator.line_end(branch.test.end()), + locator.line_end(branch.end()), + ) +} diff --git a/crates/ruff_linter/src/rules/flake8_simplify/rules/mod.rs b/crates/ruff_linter/src/rules/flake8_simplify/rules/mod.rs index 0e4b3b4937686..731eb845403dd 100644 --- a/crates/ruff_linter/src/rules/flake8_simplify/rules/mod.rs +++ b/crates/ruff_linter/src/rules/flake8_simplify/rules/mod.rs @@ -1,10 +1,15 @@ pub(crate) use ast_bool_op::*; pub(crate) use ast_expr::*; -pub(crate) use ast_if::*; pub(crate) use ast_ifexp::*; pub(crate) use ast_unary_op::*; pub(crate) use ast_with::*; +pub(crate) use collapsible_if::*; +pub(crate) use if_else_block_instead_of_dict_get::*; +pub(crate) use if_else_block_instead_of_dict_lookup::*; +pub(crate) use if_else_block_instead_of_if_exp::*; +pub(crate) use if_with_same_arms::*; pub(crate) use key_in_dict::*; +pub(crate) use needless_bool::*; pub(crate) use open_file_with_context_handler::*; pub(crate) use reimplemented_builtin::*; pub(crate) use return_in_try_except_finally::*; @@ -13,13 +18,17 @@ pub(crate) use yoda_conditions::*; mod ast_bool_op; mod ast_expr; -mod ast_if; mod ast_ifexp; mod ast_unary_op; mod ast_with; -mod fix_if; +mod collapsible_if; mod fix_with; +mod if_else_block_instead_of_dict_get; +mod if_else_block_instead_of_dict_lookup; +mod if_else_block_instead_of_if_exp; +mod if_with_same_arms; mod key_in_dict; +mod needless_bool; mod open_file_with_context_handler; mod reimplemented_builtin; mod return_in_try_except_finally; diff --git a/crates/ruff_linter/src/rules/flake8_simplify/rules/needless_bool.rs b/crates/ruff_linter/src/rules/flake8_simplify/rules/needless_bool.rs new file mode 100644 index 0000000000000..26011958a9d6d --- /dev/null +++ b/crates/ruff_linter/src/rules/flake8_simplify/rules/needless_bool.rs @@ -0,0 +1,180 @@ +use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation}; +use ruff_macros::{derive_message_formats, violation}; +use ruff_python_ast::{self as ast, Arguments, Constant, ElifElseClause, Expr, ExprContext, Stmt}; +use ruff_text_size::{Ranged, TextRange}; + +use crate::checkers::ast::Checker; +use crate::registry::AsRule; + +/// ## What it does +/// Checks for `if` statements that can be replaced with `bool`. +/// +/// ## Why is this bad? +/// `if` statements that return `True` for a truthy condition and `False` for +/// a falsey condition can be replaced with boolean casts. +/// +/// ## Example +/// ```python +/// if foo: +/// return True +/// else: +/// return False +/// ``` +/// +/// Use instead: +/// ```python +/// return bool(foo) +/// ``` +/// +/// ## References +/// - [Python documentation: Truth Value Testing](https://docs.python.org/3/library/stdtypes.html#truth-value-testing) +#[violation] +pub struct NeedlessBool { + condition: String, +} + +impl Violation for NeedlessBool { + const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes; + + #[derive_message_formats] + fn message(&self) -> String { + let NeedlessBool { condition } = self; + format!("Return the condition `{condition}` directly") + } + + fn fix_title(&self) -> Option { + let NeedlessBool { condition } = self; + Some(format!("Replace with `return {condition}`")) + } +} + +/// SIM103 +pub(crate) fn needless_bool(checker: &mut Checker, stmt: &Stmt) { + let Stmt::If(ast::StmtIf { + test: if_test, + body: if_body, + elif_else_clauses, + range: _, + }) = stmt + else { + return; + }; + // Extract an `if` or `elif` (that returns) followed by an else (that returns the same value) + let (if_test, if_body, else_body, range) = match elif_else_clauses.as_slice() { + // if-else case + [ElifElseClause { + body: else_body, + test: None, + .. + }] => (if_test.as_ref(), if_body, else_body, stmt.range()), + // elif-else case + [.., ElifElseClause { + body: elif_body, + test: Some(elif_test), + range: elif_range, + }, ElifElseClause { + body: else_body, + test: None, + range: else_range, + }] => ( + elif_test, + elif_body, + else_body, + TextRange::new(elif_range.start(), else_range.end()), + ), + _ => return, + }; + + let (Some(if_return), Some(else_return)) = ( + is_one_line_return_bool(if_body), + is_one_line_return_bool(else_body), + ) else { + return; + }; + + // If the branches have the same condition, abort (although the code could be + // simplified). + if if_return == else_return { + return; + } + + let condition = checker.generator().expr(if_test); + let mut diagnostic = Diagnostic::new(NeedlessBool { condition }, range); + if checker.patch(diagnostic.kind.rule()) { + if matches!(if_return, Bool::True) + && matches!(else_return, Bool::False) + && !checker.indexer().has_comments(&range, checker.locator()) + && (if_test.is_compare_expr() || checker.semantic().is_builtin("bool")) + { + if if_test.is_compare_expr() { + // If the condition is a comparison, we can replace it with the condition. + let node = ast::StmtReturn { + value: Some(Box::new(if_test.clone())), + range: TextRange::default(), + }; + diagnostic.set_fix(Fix::suggested(Edit::range_replacement( + checker.generator().stmt(&node.into()), + range, + ))); + } else { + // Otherwise, we need to wrap the condition in a call to `bool`. (We've already + // verified, above, that `bool` is a builtin.) + let node = ast::ExprName { + id: "bool".into(), + ctx: ExprContext::Load, + range: TextRange::default(), + }; + let node1 = ast::ExprCall { + func: Box::new(node.into()), + arguments: Arguments { + args: vec![if_test.clone()], + keywords: vec![], + range: TextRange::default(), + }, + range: TextRange::default(), + }; + let node2 = ast::StmtReturn { + value: Some(Box::new(node1.into())), + range: TextRange::default(), + }; + diagnostic.set_fix(Fix::suggested(Edit::range_replacement( + checker.generator().stmt(&node2.into()), + range, + ))); + }; + } + } + checker.diagnostics.push(diagnostic); +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Bool { + True, + False, +} + +impl From for Bool { + fn from(value: bool) -> Self { + if value { + Bool::True + } else { + Bool::False + } + } +} + +fn is_one_line_return_bool(stmts: &[Stmt]) -> Option { + let [stmt] = stmts else { + return None; + }; + let Stmt::Return(ast::StmtReturn { value, range: _ }) = stmt else { + return None; + }; + let Some(Expr::Constant(ast::ExprConstant { value, .. })) = value.as_deref() else { + return None; + }; + let Constant::Bool(value) = value else { + return None; + }; + Some((*value).into()) +}