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
```

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));
                 }
             }
         }

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<Loop> {
+    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<Loop<'a>> {
+fn return_values_for_siblings<'a>(stmt: &'a Stmt, sibling: &'a Stmt) -> Option<Loop<'a>> {
     let StmtKind::For {
         body,
         target,
@@ -36,14 +99,13 @@ fn return_values<'a>(stmt: &'a Stmt, sibling: &'a Stmt) -> Option<Loop<'a>> {
     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);

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,

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: ~
 

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: ~