From 06473bb1b5408affbb70ccbd98cc3a4730eec031 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Thu, 12 Jan 2023 17:04:58 -0500 Subject: [PATCH] Support for-else loops in `SIM110` and `SIM111` (#1834) This PR adds support for `SIM110` and `SIM111` simplifications of the form: ```py def f(): # SIM110 for x in iterable: if check(x): return True else: return False ``` --- .../test/fixtures/flake8_simplify/SIM110.py | 38 ++++++++ .../test/fixtures/flake8_simplify/SIM111.py | 38 ++++++++ src/checkers/ast.rs | 13 ++- src/flake8_simplify/rules/ast_for.rs | 91 +++++++++++++++++-- src/flake8_simplify/rules/mod.rs | 2 +- ...ke8_simplify__tests__SIM110_SIM110.py.snap | 34 +++++++ ...ke8_simplify__tests__SIM111_SIM111.py.snap | 34 +++++++ 7 files changed, 237 insertions(+), 13 deletions(-) diff --git a/resources/test/fixtures/flake8_simplify/SIM110.py b/resources/test/fixtures/flake8_simplify/SIM110.py index dbf3e64420..ddf6f82d6c 100644 --- a/resources/test/fixtures/flake8_simplify/SIM110.py +++ b/resources/test/fixtures/flake8_simplify/SIM110.py @@ -50,6 +50,44 @@ def f(): return "bar" +def f(): + # SIM110 + for x in iterable: + if check(x): + return True + else: + return False + + +def f(): + # SIM111 + for x in iterable: + if check(x): + return False + else: + return True + + +def f(): + # SIM110 + for x in iterable: + if check(x): + return True + else: + return False + return True + + +def f(): + # SIM111 + for x in iterable: + if check(x): + return False + else: + return True + return False + + def f(): for x in iterable: if check(x): diff --git a/resources/test/fixtures/flake8_simplify/SIM111.py b/resources/test/fixtures/flake8_simplify/SIM111.py index dbf3e64420..ddf6f82d6c 100644 --- a/resources/test/fixtures/flake8_simplify/SIM111.py +++ b/resources/test/fixtures/flake8_simplify/SIM111.py @@ -50,6 +50,44 @@ def f(): return "bar" +def f(): + # SIM110 + for x in iterable: + if check(x): + return True + else: + return False + + +def f(): + # SIM111 + for x in iterable: + if check(x): + return False + else: + return True + + +def f(): + # SIM110 + for x in iterable: + if check(x): + return True + else: + return False + return True + + +def f(): + # SIM111 + for x in iterable: + if check(x): + return False + else: + return True + return False + + def f(): for x in iterable: if check(x): diff --git a/src/checkers/ast.rs b/src/checkers/ast.rs index cbba233d94..5eaddc4bad 100644 --- a/src/checkers/ast.rs +++ b/src/checkers/ast.rs @@ -1310,8 +1310,15 @@ where if self.settings.enabled.contains(&RuleCode::PLW0120) { pylint::rules::useless_else_on_loop(self, stmt, body, orelse); } - if self.settings.enabled.contains(&RuleCode::SIM118) { - flake8_simplify::rules::key_in_dict_for(self, target, iter); + if matches!(stmt.node, StmtKind::For { .. }) { + if self.settings.enabled.contains(&RuleCode::SIM110) + || self.settings.enabled.contains(&RuleCode::SIM111) + { + flake8_simplify::rules::convert_for_loop_to_any_all(self, stmt, None); + } + if self.settings.enabled.contains(&RuleCode::SIM118) { + flake8_simplify::rules::key_in_dict_for(self, target, iter); + } } } StmtKind::Try { @@ -3168,7 +3175,7 @@ where if matches!(stmt.node, StmtKind::For { .. }) && matches!(sibling.node, StmtKind::Return { .. }) { - flake8_simplify::rules::convert_loop_to_any_all(self, stmt, sibling); + flake8_simplify::rules::convert_for_loop_to_any_all(self, stmt, Some(sibling)); } } } diff --git a/src/flake8_simplify/rules/ast_for.rs b/src/flake8_simplify/rules/ast_for.rs index e10c7c84e2..57ae21bb82 100644 --- a/src/flake8_simplify/rules/ast_for.rs +++ b/src/flake8_simplify/rules/ast_for.rs @@ -18,9 +18,72 @@ struct Loop<'a> { iter: &'a Expr, } -/// Extract the returned boolean values from subsequent `StmtKind::If` and +/// Extract the returned boolean values a `StmtKind::For` with an `else` body. +fn return_values_for_else(stmt: &Stmt) -> Option { + let StmtKind::For { + body, + target, + iter, + orelse, + .. + } = &stmt.node else { + return None; + }; + + // The loop itself should contain a single `if` statement, with an `else` + // containing a single `return True` or `return False`. + if body.len() != 1 { + return None; + } + if orelse.len() != 1 { + return None; + } + let StmtKind::If { + body: nested_body, + test: nested_test, + orelse: nested_orelse, + } = &body[0].node else { + return None; + }; + if nested_body.len() != 1 { + return None; + } + if !nested_orelse.is_empty() { + return None; + } + let StmtKind::Return { value } = &nested_body[0].node else { + return None; + }; + let Some(value) = value else { + return None; + }; + let ExprKind::Constant { value: Constant::Bool(value), .. } = &value.node else { + return None; + }; + + // The `else` block has to contain a single `return True` or `return False`. + let StmtKind::Return { value: next_value } = &orelse[0].node else { + return None; + }; + let Some(next_value) = next_value else { + return None; + }; + let ExprKind::Constant { value: Constant::Bool(next_value), .. } = &next_value.node else { + return None; + }; + + Some(Loop { + return_value: *value, + next_return_value: *next_value, + test: nested_test, + target, + iter, + }) +} + +/// Extract the returned boolean values from subsequent `StmtKind::For` and /// `StmtKind::Return` statements, or `None`. -fn return_values<'a>(stmt: &'a Stmt, sibling: &'a Stmt) -> Option> { +fn return_values_for_siblings<'a>(stmt: &'a Stmt, sibling: &'a Stmt) -> Option> { let StmtKind::For { body, target, @@ -36,14 +99,13 @@ fn return_values<'a>(stmt: &'a Stmt, sibling: &'a Stmt) -> Option> { if body.len() != 1 { return None; } - // TODO(charlie): If we have `else: return True` or `else: return False`, we - // should still be able to simplify. if !orelse.is_empty() { return None; } let StmtKind::If { body: nested_body, - test: nested_test, orelse: nested_orelse, + test: nested_test, + orelse: nested_orelse, } = &body[0].node else { return None; }; @@ -108,8 +170,13 @@ fn return_stmt(id: &str, test: &Expr, target: &Expr, iter: &Expr, stylist: &Styl } /// SIM110, SIM111 -pub fn convert_loop_to_any_all(checker: &mut Checker, stmt: &Stmt, sibling: &Stmt) { - if let Some(loop_info) = return_values(stmt, sibling) { +pub fn convert_for_loop_to_any_all(checker: &mut Checker, stmt: &Stmt, sibling: Option<&Stmt>) { + if let Some(loop_info) = match sibling { + // Ex) `for` loop with an `else: return True` or `else: return False`. + None => return_values_for_else(stmt), + // Ex) `for` loop followed by `return True` or `return False` + Some(sibling) => return_values_for_siblings(stmt, sibling), + } { if loop_info.return_value && !loop_info.next_return_value { if checker.settings.enabled.contains(&RuleCode::SIM110) { let contents = return_stmt( @@ -133,7 +200,10 @@ pub fn convert_loop_to_any_all(checker: &mut Checker, stmt: &Stmt, sibling: &Stm diagnostic.amend(Fix::replacement( contents, stmt.location, - sibling.end_location.unwrap(), + match sibling { + None => stmt.end_location.unwrap(), + Some(sibling) => sibling.end_location.unwrap(), + }, )); } checker.diagnostics.push(diagnostic); @@ -178,7 +248,10 @@ pub fn convert_loop_to_any_all(checker: &mut Checker, stmt: &Stmt, sibling: &Stm diagnostic.amend(Fix::replacement( contents, stmt.location, - sibling.end_location.unwrap(), + match sibling { + None => stmt.end_location.unwrap(), + Some(sibling) => sibling.end_location.unwrap(), + }, )); } checker.diagnostics.push(diagnostic); diff --git a/src/flake8_simplify/rules/mod.rs b/src/flake8_simplify/rules/mod.rs index cbb5d4de47..b50dc07257 100644 --- a/src/flake8_simplify/rules/mod.rs +++ b/src/flake8_simplify/rules/mod.rs @@ -2,7 +2,7 @@ pub use ast_bool_op::{ a_and_not_a, a_or_not_a, and_false, compare_with_tuple, duplicate_isinstance_call, or_true, }; pub use ast_expr::use_capital_environment_variables; -pub use ast_for::convert_loop_to_any_all; +pub use ast_for::convert_for_loop_to_any_all; pub use ast_if::{ nested_if_statements, return_bool_condition_directly, use_dict_get_with_default, use_ternary_operator, diff --git a/src/flake8_simplify/snapshots/ruff__flake8_simplify__tests__SIM110_SIM110.py.snap b/src/flake8_simplify/snapshots/ruff__flake8_simplify__tests__SIM110_SIM110.py.snap index 05cc0c3996..8f54cc3c71 100644 --- a/src/flake8_simplify/snapshots/ruff__flake8_simplify__tests__SIM110_SIM110.py.snap +++ b/src/flake8_simplify/snapshots/ruff__flake8_simplify__tests__SIM110_SIM110.py.snap @@ -19,4 +19,38 @@ expression: diagnostics row: 6 column: 16 parent: ~ +- kind: + ConvertLoopToAny: return any(check(x) for x in iterable) + location: + row: 55 + column: 4 + end_location: + row: 59 + column: 20 + fix: + content: return any(check(x) for x in iterable) + location: + row: 55 + column: 4 + end_location: + row: 59 + column: 20 + parent: ~ +- kind: + ConvertLoopToAny: return any(check(x) for x in iterable) + location: + row: 73 + column: 4 + end_location: + row: 77 + column: 20 + fix: + content: return any(check(x) for x in iterable) + location: + row: 73 + column: 4 + end_location: + row: 77 + column: 20 + parent: ~ diff --git a/src/flake8_simplify/snapshots/ruff__flake8_simplify__tests__SIM111_SIM111.py.snap b/src/flake8_simplify/snapshots/ruff__flake8_simplify__tests__SIM111_SIM111.py.snap index 95fc7e6cfa..2443a844c3 100644 --- a/src/flake8_simplify/snapshots/ruff__flake8_simplify__tests__SIM111_SIM111.py.snap +++ b/src/flake8_simplify/snapshots/ruff__flake8_simplify__tests__SIM111_SIM111.py.snap @@ -36,4 +36,38 @@ expression: diagnostics row: 36 column: 15 parent: ~ +- kind: + ConvertLoopToAll: return all(not check(x) for x in iterable) + location: + row: 64 + column: 4 + end_location: + row: 68 + column: 19 + fix: + content: return all(not check(x) for x in iterable) + location: + row: 64 + column: 4 + end_location: + row: 68 + column: 19 + parent: ~ +- kind: + ConvertLoopToAll: return all(not check(x) for x in iterable) + location: + row: 83 + column: 4 + end_location: + row: 87 + column: 19 + fix: + content: return all(not check(x) for x in iterable) + location: + row: 83 + column: 4 + end_location: + row: 87 + column: 19 + parent: ~