Implement SIM110 and SIM111 (conversion to any and all) (#1653)

README.md

@@ -965,10 +965,12 @@ For more, see [flake8-simplify](https://pypi.org/project/flake8-simplify/0.19.3/
 
 | Code | Name | Message | Fix |
 | ---- | ---- | ------- | --- |
-| SIM117 | MultipleWithStatements | Use a single `with` statement with multiple contexts instead of nested `with` statements |  |
 | SIM102 | NestedIfStatements | Use a single `if` statement instead of nested `if` statements |  |
 | SIM105 | UseContextlibSuppress | Use `contextlib.suppress(...)` instead of try-except-pass |  |
 | SIM107 | ReturnInTryExceptFinally | Don't use `return` in `try`/`except` and `finally` |  |
+| SIM110 | ConvertLoopToAny | Use `return any(x for x in y)` instead of `for` loop | 🛠 |
+| SIM111 | ConvertLoopToAll | Use `return all(x for x in y)` instead of `for` loop | 🛠 |
+| SIM117 | MultipleWithStatements | Use a single `with` statement with multiple contexts instead of nested `with` statements |  |
 | SIM118 | KeyInDict | Use `key in dict` instead of `key in dict.keys()` | 🛠 |
 | SIM220 | AAndNotA | Use `False` instead of `... and not ...` | 🛠 |
 | SIM221 | AOrNotA | Use `True` instead of `... or not ...` | 🛠 |

resources/test/fixtures/flake8_simplify/SIM110.py

@@ -0,0 +1,47 @@
+def f():
+    for x in iterable:  # SIM110
+        if check(x):
+            return True
+    return False
+
+
+def f():
+    for x in iterable:
+        if check(x):
+            return True
+    return True
+
+
+def f():
+    for el in [1, 2, 3]:
+        if is_true(el):
+            return True
+    raise Exception
+
+
+def f():
+    for x in iterable:  # SIM111
+        if check(x):
+            return False
+    return True
+
+
+def f():
+    for x in iterable:  # SIM111
+        if not x.is_empty():
+            return False
+    return True
+
+
+def f():
+    for x in iterable:
+        if check(x):
+            return False
+    return False
+
+
+def f():
+    for x in iterable:
+        if check(x):
+            return "foo"
+    return "bar"

resources/test/fixtures/flake8_simplify/SIM111.py

@@ -0,0 +1,33 @@
+def f():
+    for x in iterable:  # SIM110
+        if check(x):
+            return True
+    return False
+
+
+def f():
+    for el in [1, 2, 3]:
+        if is_true(el):
+            return True
+    raise Exception
+
+
+def f():
+    for x in iterable:  # SIM111
+        if check(x):
+            return False
+    return True
+
+
+def f():
+    for x in iterable:  # SIM 111
+        if not x.is_empty():
+            return False
+    return True
+
+
+def f():
+    for x in iterable:
+        if check(x):
+            return "foo"
+    return "bar"

src/checkers/ast.rs

@@ -3061,6 +3061,19 @@ where
         if self.settings.enabled.contains(&CheckCode::PIE790) {
             flake8_pie::plugins::no_unnecessary_pass(self, body);
         }
+
+        if self.settings.enabled.contains(&CheckCode::SIM110)
+            || self.settings.enabled.contains(&CheckCode::SIM111)
+        {
+            for (stmt, sibling) in body.iter().tuple_windows() {
+                if matches!(stmt.node, StmtKind::For { .. })
+                    && matches!(sibling.node, StmtKind::Return { .. })
+                {
+                    flake8_simplify::plugins::convert_loop_to_any_all(self, stmt, sibling);
+                }
+            }
+        }
+
         visitor::walk_body(self, body);
     }
 }

src/flake8_simplify/mod.rs

@@ -15,6 +15,8 @@ mod tests {
     #[test_case(CheckCode::SIM102, Path::new("SIM102.py"); "SIM102")]
     #[test_case(CheckCode::SIM105, Path::new("SIM105.py"); "SIM105")]
     #[test_case(CheckCode::SIM107, Path::new("SIM107.py"); "SIM107")]
+    #[test_case(CheckCode::SIM110, Path::new("SIM110.py"); "SIM110")]
+    #[test_case(CheckCode::SIM111, Path::new("SIM111.py"); "SIM111")]
     #[test_case(CheckCode::SIM117, Path::new("SIM117.py"); "SIM117")]
     #[test_case(CheckCode::SIM118, Path::new("SIM118.py"); "SIM118")]
     #[test_case(CheckCode::SIM220, Path::new("SIM220.py"); "SIM220")]

src/flake8_simplify/plugins/ast_for.rs

@@ -0,0 +1,187 @@
+use log::error;
+use rustpython_ast::{
+    Comprehension, Constant, Expr, ExprContext, ExprKind, Stmt, StmtKind, Unaryop,
+};
+
+use crate::ast::helpers::{create_expr, create_stmt};
+use crate::ast::types::Range;
+use crate::autofix::Fix;
+use crate::checkers::ast::Checker;
+use crate::registry::{Check, CheckCode, CheckKind};
+use crate::source_code_generator::SourceCodeGenerator;
+use crate::source_code_style::SourceCodeStyleDetector;
+
+struct Loop<'a> {
+    return_value: bool,
+    next_return_value: bool,
+    test: &'a Expr,
+    target: &'a Expr,
+    iter: &'a Expr,
+}
+
+/// Extract the returned boolean values from subsequent `StmtKind::If` and
+/// `StmtKind::Return` statements, or `None`.
+fn return_values<'a>(stmt: &'a Stmt, sibling: &'a Stmt) -> Option<Loop<'a>> {
+    let StmtKind::For {
+        body,
+        target,
+        iter,
+        ..
+    } = &stmt.node else {
+        return None;
+    };
+
+    // The loop itself should contain a single `if` statement, with a single `return
+    // True` or `return False`.
+    if body.len() != 1 {
+        return None;
+    }
+    let StmtKind::If {
+        body: nested_body,
+        test: nested_test,
+        ..
+    } = &body[0].node else {
+        return None;
+    };
+    if nested_body.len() != 1 {
+        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 next statement has to be a `return True` or `return False`.
+    let StmtKind::Return { value: next_value } = &sibling.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,
+    })
+}
+
+/// Generate a return statement for an `any` or `all` builtin comprehension.
+fn return_stmt(
+    id: &str,
+    test: &Expr,
+    target: &Expr,
+    iter: &Expr,
+    stylist: &SourceCodeStyleDetector,
+) -> Option<String> {
+    let mut generator = SourceCodeGenerator::new(
+        stylist.indentation(),
+        stylist.quote(),
+        stylist.line_ending(),
+    );
+    generator.unparse_stmt(&create_stmt(StmtKind::Return {
+        value: Some(Box::new(create_expr(ExprKind::Call {
+            func: Box::new(create_expr(ExprKind::Name {
+                id: id.to_string(),
+                ctx: ExprContext::Load,
+            })),
+            args: vec![create_expr(ExprKind::GeneratorExp {
+                elt: Box::new(test.clone()),
+                generators: vec![Comprehension {
+                    target: target.clone(),
+                    iter: iter.clone(),
+                    ifs: vec![],
+                    is_async: 0,
+                }],
+            })],
+            keywords: vec![],
+        }))),
+    }));
+    match generator.generate() {
+        Ok(test) => Some(test),
+        Err(e) => {
+            error!("Failed to generate source code: {}", e);
+            None
+        }
+    }
+}
+
+/// 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) {
+        if loop_info.return_value && !loop_info.next_return_value {
+            if checker.settings.enabled.contains(&CheckCode::SIM110) {
+                if let Some(content) = return_stmt(
+                    "any",
+                    loop_info.test,
+                    loop_info.target,
+                    loop_info.iter,
+                    checker.style,
+                ) {
+                    let mut check = Check::new(
+                        CheckKind::ConvertLoopToAny(content.clone()),
+                        Range::from_located(stmt),
+                    );
+                    if checker.patch(&CheckCode::SIM110) {
+                        check.amend(Fix::replacement(
+                            content,
+                            stmt.location,
+                            sibling.end_location.unwrap(),
+                        ));
+                    }
+                    checker.add_check(check);
+                }
+            }
+        }
+
+        if !loop_info.return_value && loop_info.next_return_value {
+            if checker.settings.enabled.contains(&CheckCode::SIM111) {
+                // Invert the condition.
+                let test = {
+                    if let ExprKind::UnaryOp {
+                        op: Unaryop::Not,
+                        operand,
+                    } = &loop_info.test.node
+                    {
+                        *operand.clone()
+                    } else {
+                        create_expr(ExprKind::UnaryOp {
+                            op: Unaryop::Not,
+                            operand: Box::new(loop_info.test.clone()),
+                        })
+                    }
+                };
+                if let Some(content) = return_stmt(
+                    "all",
+                    &test,
+                    loop_info.target,
+                    loop_info.iter,
+                    checker.style,
+                ) {
+                    let mut check = Check::new(
+                        CheckKind::ConvertLoopToAll(content.clone()),
+                        Range::from_located(stmt),
+                    );
+                    if checker.patch(&CheckCode::SIM111) {
+                        check.amend(Fix::replacement(
+                            content,
+                            stmt.location,
+                            sibling.end_location.unwrap(),
+                        ));
+                    }
+                    checker.add_check(check);
+                }
+            }
+        }
+    }
+}

src/flake8_simplify/plugins/mod.rs

@@ -1,4 +1,5 @@
 pub use ast_bool_op::{a_and_not_a, a_or_not_a, and_false, or_true};
+pub use ast_for::convert_loop_to_any_all;
 pub use ast_if::nested_if_statements;
 pub use ast_with::multiple_with_statements;
 pub use key_in_dict::{key_in_dict_compare, key_in_dict_for};
@@ -7,6 +8,7 @@ pub use use_contextlib_suppress::use_contextlib_suppress;
 pub use yoda_conditions::yoda_conditions;
 
 mod ast_bool_op;
+mod ast_for;
 mod ast_if;
 mod ast_with;
 mod key_in_dict;

src/flake8_simplify/snapshots/ruff__flake8_simplify__tests__SIM110_SIM110.py.snap

@@ -0,0 +1,22 @@
+---
+source: src/flake8_simplify/mod.rs
+expression: checks
+---
+- kind:
+    ConvertLoopToAny: return any(check(x) for x in iterable)
+  location:
+    row: 2
+    column: 4
+  end_location:
+    row: 4
+    column: 23
+  fix:
+    content: return any(check(x) for x in iterable)
+    location:
+      row: 2
+      column: 4
+    end_location:
+      row: 5
+      column: 16
+  parent: ~
+

src/flake8_simplify/snapshots/ruff__flake8_simplify__tests__SIM111_SIM111.py.snap

@@ -0,0 +1,39 @@
+---
+source: src/flake8_simplify/mod.rs
+expression: checks
+---
+- kind:
+    ConvertLoopToAll: return all(not check(x) for x in iterable)
+  location:
+    row: 16
+    column: 4
+  end_location:
+    row: 18
+    column: 24
+  fix:
+    content: return all(not check(x) for x in iterable)
+    location:
+      row: 16
+      column: 4
+    end_location:
+      row: 19
+      column: 15
+  parent: ~
+- kind:
+    ConvertLoopToAll: return all(x.is_empty() for x in iterable)
+  location:
+    row: 23
+    column: 4
+  end_location:
+    row: 25
+    column: 24
+  fix:
+    content: return all(x.is_empty() for x in iterable)
+    location:
+      row: 23
+      column: 4
+    end_location:
+      row: 26
+      column: 15
+  parent: ~
+