From 59c00b5298d3272077120695ba3aa3ecd6ee6842 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 4 Oct 2023 14:18:59 -0400 Subject: [PATCH] Use a dedicated struct for "nested if" rule (#7817) Internal refactor -- finding this rule hard to understand. --- .../src/rules/flake8_simplify/rules/ast_if.rs | 94 +++++++++++++------ .../src/rules/flake8_simplify/rules/fix_if.rs | 18 ++-- 2 files changed, 73 insertions(+), 39 deletions(-) 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 index ec76d2c4e3..bb9d5c19d3 100644 --- a/crates/ruff_linter/src/rules/flake8_simplify/rules/ast_if.rs +++ b/crates/ruff_linter/src/rules/flake8_simplify/rules/ast_if.rs @@ -1,18 +1,19 @@ use log::error; -use ruff_python_ast::{ - self as ast, Arguments, CmpOp, Constant, ElifElseClause, Expr, ExprContext, Identifier, Stmt, -}; -use ruff_text_size::{Ranged, TextRange}; 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, UniversalNewlines}; +use ruff_text_size::{Ranged, TextRange}; use crate::checkers::ast::Checker; use crate::line_width::LineWidthBuilder; @@ -286,10 +287,9 @@ fn is_main_check(expr: &Expr) -> bool { /// if yyy: /// # ^^^ returns this expression /// z = 1 -/// # ^^^^^ and this statement /// ... /// ``` -fn find_last_nested_if(body: &[Stmt]) -> Option<(&Expr, &Stmt)> { +fn find_last_nested_if(body: &[Stmt]) -> Option<&Expr> { let [Stmt::If(ast::StmtIf { test, body: inner_body, @@ -302,16 +302,48 @@ fn find_last_nested_if(body: &[Stmt]) -> Option<(&Expr, &Stmt)> { if !elif_else_clauses.is_empty() { return None; } - find_last_nested_if(inner_body).or_else(|| { - Some(( - test, - inner_body.last().expect("Expected body to be non-empty"), - )) - }) + 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<(&[Stmt], TextRange, bool)> { +fn nested_if_body(stmt_if: &ast::StmtIf) -> Option { let ast::StmtIf { test, body, @@ -321,15 +353,15 @@ fn nested_if_body(stmt_if: &ast::StmtIf) -> Option<(&[Stmt], TextRange, bool)> { // 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, body, range, is_elif) = if let Some(clause) = elif_else_clauses.last() { + let (test, nested_if) = if let Some(clause) = elif_else_clauses.last() { if let Some(test) = &clause.test { - (test, &clause.body, clause.range(), true) + (test, NestedIf::Elif(clause)) } else { // The last condition is an `else` (different rule) return None; } } else { - (test.as_ref(), body, stmt_if.range(), false) + (test.as_ref(), NestedIf::If(stmt_if)) }; // The nested if must be the only child, otherwise there is at least one more statement that @@ -354,7 +386,7 @@ fn nested_if_body(stmt_if: &ast::StmtIf) -> Option<(&[Stmt], TextRange, bool)> { return None; } - Some((body, range, is_elif)) + Some(nested_if) } /// SIM102 @@ -363,21 +395,22 @@ pub(crate) fn nested_if_statements( stmt_if: &ast::StmtIf, parent: Option<&Stmt>, ) { - let Some((body, range, is_elif)) = nested_if_body(stmt_if) else { + let Some(nested_if) = nested_if_body(stmt_if) else { return; }; // Find the deepest nested if-statement, to inform the range. - let Some((test, _first_stmt)) = find_last_nested_if(body) else { + 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((body, _range, _is_elif)) = nested_if_body(stmt_if) { + 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 - if matches!(&body[0], Stmt::If(inner) if inner == stmt_if) + let body = nested_if.body(); + if matches!(&body[0], Stmt::If(inner) if *inner == *stmt_if) && find_last_nested_if(body).is_some() { return; @@ -392,22 +425,23 @@ pub(crate) fn nested_if_statements( return; }; - let mut diagnostic = Diagnostic::new(CollapsibleIf, TextRange::new(range.start(), colon.end())); + 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. - let nested_if = &body[0]; if !checker .indexer() .comment_ranges() - .intersects(TextRange::new(range.start(), nested_if.start())) + .intersects(TextRange::new( + nested_if.start(), + nested_if.body()[0].start(), + )) { - match fix_if::fix_nested_if_statements( - checker.locator(), - checker.stylist(), - range, - is_elif, - ) { + match fix_if::fix_nested_if_statements(checker.locator(), checker.stylist(), nested_if) + { Ok(edit) => { if edit .content() 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 index f308adf710..bf775fb365 100644 --- a/crates/ruff_linter/src/rules/flake8_simplify/rules/fix_if.rs +++ b/crates/ruff_linter/src/rules/flake8_simplify/rules/fix_if.rs @@ -10,11 +10,12 @@ use ruff_diagnostics::Edit; use ruff_python_ast::whitespace; use ruff_python_codegen::Stylist; use ruff_source_file::Locator; -use ruff_text_size::TextRange; +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 { @@ -32,23 +33,22 @@ fn parenthesize_and_operand(expr: Expression) -> Expression { } /// (SIM102) Convert `if a: if b:` to `if a and b:`. -pub(crate) fn fix_nested_if_statements( +pub(super) fn fix_nested_if_statements( locator: &Locator, stylist: &Stylist, - range: TextRange, - is_elif: bool, + nested_if: NestedIf, ) -> Result { // Infer the indentation of the outer block. - let Some(outer_indent) = whitespace::indentation(locator, &range) else { + let Some(outer_indent) = whitespace::indentation(locator, &nested_if) else { bail!("Unable to fix multiline statement"); }; // Extract the module text. - let contents = locator.lines(range); + 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 is_elif { + let module_text = if nested_if.is_elif() { Cow::Owned(contents.replacen("elif", "if", 1)) } else { Cow::Borrowed(contents) @@ -121,12 +121,12 @@ pub(crate) fn fix_nested_if_statements( .strip_prefix(&format!("def f():{}", stylist.line_ending().as_str())) .unwrap() }; - let contents = if is_elif { + 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(range); + let range = locator.lines_range(nested_if.range()); Ok(Edit::range_replacement(contents.to_string(), range)) }