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

src/registry.rs

@@ -217,10 +217,12 @@ pub enum CheckCode {
     YTT302,
     YTT303,
     // flake8-simplify
-    SIM117,
     SIM102,
     SIM105,
     SIM107,
+    SIM110,
+    SIM111,
+    SIM117,
     SIM118,
     SIM220,
     SIM221,
@@ -954,15 +956,17 @@ pub enum CheckKind {
     SysVersionCmpStr10,
     SysVersionSlice1Referenced,
     // flake8-simplify
-    MultipleWithStatements,
-    NestedIfStatements,
-    UseContextlibSuppress(String),
-    ReturnInTryExceptFinally,
     AAndNotA(String),
     AOrNotA(String),
-    KeyInDict(String, String),
     AndFalse,
+    ConvertLoopToAll(String),
+    ConvertLoopToAny(String),
+    KeyInDict(String, String),
+    MultipleWithStatements,
+    NestedIfStatements,
     OrTrue,
+    ReturnInTryExceptFinally,
+    UseContextlibSuppress(String),
     YodaConditions(String, String),
     // pyupgrade
     ConvertNamedTupleFunctionalToClass(String),
@@ -1394,6 +1398,12 @@ impl CheckCode {
             CheckCode::SIM102 => CheckKind::NestedIfStatements,
             CheckCode::SIM105 => CheckKind::UseContextlibSuppress("...".to_string()),
             CheckCode::SIM107 => CheckKind::ReturnInTryExceptFinally,
+            CheckCode::SIM110 => {
+                CheckKind::ConvertLoopToAny("return any(x for x in y)".to_string())
+            }
+            CheckCode::SIM111 => {
+                CheckKind::ConvertLoopToAll("return all(x for x in y)".to_string())
+            }
             CheckCode::SIM117 => CheckKind::MultipleWithStatements,
             CheckCode::SIM118 => CheckKind::KeyInDict("key".to_string(), "dict".to_string()),
             CheckCode::SIM220 => CheckKind::AAndNotA("...".to_string()),
@@ -1930,6 +1940,8 @@ impl CheckCode {
             CheckCode::SIM102 => CheckCategory::Flake8Simplify,
             CheckCode::SIM105 => CheckCategory::Flake8Simplify,
             CheckCode::SIM107 => CheckCategory::Flake8Simplify,
+            CheckCode::SIM110 => CheckCategory::Flake8Simplify,
+            CheckCode::SIM111 => CheckCategory::Flake8Simplify,
             CheckCode::SIM117 => CheckCategory::Flake8Simplify,
             CheckCode::SIM118 => CheckCategory::Flake8Simplify,
             CheckCode::SIM220 => CheckCategory::Flake8Simplify,
@@ -2183,15 +2195,17 @@ impl CheckKind {
             CheckKind::SysVersionCmpStr10 => &CheckCode::YTT302,
             CheckKind::SysVersionSlice1Referenced => &CheckCode::YTT303,
             // flake8-simplify
-            CheckKind::NestedIfStatements => &CheckCode::SIM102,
-            CheckKind::UseContextlibSuppress(..) => &CheckCode::SIM105,
-            CheckKind::ReturnInTryExceptFinally => &CheckCode::SIM107,
-            CheckKind::MultipleWithStatements => &CheckCode::SIM117,
-            CheckKind::KeyInDict(..) => &CheckCode::SIM118,
             CheckKind::AAndNotA(..) => &CheckCode::SIM220,
             CheckKind::AOrNotA(..) => &CheckCode::SIM221,
-            CheckKind::OrTrue => &CheckCode::SIM222,
             CheckKind::AndFalse => &CheckCode::SIM223,
+            CheckKind::ConvertLoopToAll(..) => &CheckCode::SIM111,
+            CheckKind::ConvertLoopToAny(..) => &CheckCode::SIM110,
+            CheckKind::KeyInDict(..) => &CheckCode::SIM118,
+            CheckKind::MultipleWithStatements => &CheckCode::SIM117,
+            CheckKind::NestedIfStatements => &CheckCode::SIM102,
+            CheckKind::OrTrue => &CheckCode::SIM222,
+            CheckKind::ReturnInTryExceptFinally => &CheckCode::SIM107,
+            CheckKind::UseContextlibSuppress(..) => &CheckCode::SIM105,
             CheckKind::YodaConditions(..) => &CheckCode::SIM300,
             // pyupgrade
             CheckKind::ConvertNamedTupleFunctionalToClass(..) => &CheckCode::UP014,
@@ -2956,6 +2970,12 @@ impl CheckKind {
             CheckKind::AAndNotA(name) => format!("Use `False` instead of `{name} and not {name}`"),
             CheckKind::AOrNotA(name) => format!("Use `True` instead of `{name} or not {name}`"),
             CheckKind::AndFalse => "Use `False` instead of `... and False`".to_string(),
+            CheckKind::ConvertLoopToAll(all) => {
+                format!("Use `{all}` instead of `for` loop")
+            }
+            CheckKind::ConvertLoopToAny(any) => {
+                format!("Use `{any}` instead of `for` loop")
+            }
             CheckKind::KeyInDict(key, dict) => {
                 format!("Use `{key} in {dict}` instead of `{key} in {dict}.keys()`")
             }
@@ -3585,6 +3605,8 @@ impl CheckKind {
                 | CheckKind::BlankLineBeforeSection(..)
                 | CheckKind::CapitalizeSectionName(..)
                 | CheckKind::CommentedOutCode
+                | CheckKind::ConvertLoopToAll(..)
+                | CheckKind::ConvertLoopToAny(..)
                 | CheckKind::ConvertNamedTupleFunctionalToClass(..)
                 | CheckKind::ConvertTypedDictFunctionalToClass(..)
                 | CheckKind::DashedUnderlineAfterSection(..)
@@ -3712,6 +3734,8 @@ impl CheckKind {
             }
             CheckKind::CapitalizeSectionName(name) => Some(format!("Capitalize \"{name}\"")),
             CheckKind::CommentedOutCode => Some("Remove commented-out code".to_string()),
+            CheckKind::ConvertLoopToAll(all) => Some(format!("Replace with `{all}`")),
+            CheckKind::ConvertLoopToAny(any) => Some(format!("Replace with `{any}`")),
             CheckKind::ConvertTypedDictFunctionalToClass(name)
             | CheckKind::ConvertNamedTupleFunctionalToClass(name) => {
                 Some(format!("Convert `{name}` to class syntax"))

src/registry_gen.rs

@@ -522,6 +522,8 @@ pub enum CheckCodePrefix {
     SIM105,
     SIM107,
     SIM11,
+    SIM110,
+    SIM111,
     SIM117,
     SIM118,
     SIM2,
@@ -809,10 +811,12 @@ impl CheckCodePrefix {
                 CheckCode::YTT301,
                 CheckCode::YTT302,
                 CheckCode::YTT303,
-                CheckCode::SIM117,
                 CheckCode::SIM102,
                 CheckCode::SIM105,
                 CheckCode::SIM107,
+                CheckCode::SIM110,
+                CheckCode::SIM111,
+                CheckCode::SIM117,
                 CheckCode::SIM118,
                 CheckCode::SIM220,
                 CheckCode::SIM221,
@@ -2635,10 +2639,12 @@ impl CheckCodePrefix {
             CheckCodePrefix::S107 => vec![CheckCode::S107],
             CheckCodePrefix::S108 => vec![CheckCode::S108],
             CheckCodePrefix::SIM => vec![
-                CheckCode::SIM117,
                 CheckCode::SIM102,
                 CheckCode::SIM105,
                 CheckCode::SIM107,
+                CheckCode::SIM110,
+                CheckCode::SIM111,
+                CheckCode::SIM117,
                 CheckCode::SIM118,
                 CheckCode::SIM220,
                 CheckCode::SIM221,
@@ -2647,17 +2653,26 @@ impl CheckCodePrefix {
                 CheckCode::SIM300,
             ],
             CheckCodePrefix::SIM1 => vec![
-                CheckCode::SIM117,
                 CheckCode::SIM102,
                 CheckCode::SIM105,
                 CheckCode::SIM107,
+                CheckCode::SIM110,
+                CheckCode::SIM111,
+                CheckCode::SIM117,
                 CheckCode::SIM118,
             ],
             CheckCodePrefix::SIM10 => vec![CheckCode::SIM102, CheckCode::SIM105, CheckCode::SIM107],
             CheckCodePrefix::SIM102 => vec![CheckCode::SIM102],
             CheckCodePrefix::SIM105 => vec![CheckCode::SIM105],
             CheckCodePrefix::SIM107 => vec![CheckCode::SIM107],
-            CheckCodePrefix::SIM11 => vec![CheckCode::SIM117, CheckCode::SIM118],
+            CheckCodePrefix::SIM11 => vec![
+                CheckCode::SIM110,
+                CheckCode::SIM111,
+                CheckCode::SIM117,
+                CheckCode::SIM118,
+            ],
+            CheckCodePrefix::SIM110 => vec![CheckCode::SIM110],
+            CheckCodePrefix::SIM111 => vec![CheckCode::SIM111],
             CheckCodePrefix::SIM117 => vec![CheckCode::SIM117],
             CheckCodePrefix::SIM118 => vec![CheckCode::SIM118],
             CheckCodePrefix::SIM2 => vec![
@@ -3634,6 +3649,8 @@ impl CheckCodePrefix {
             CheckCodePrefix::SIM105 => SuffixLength::Three,
             CheckCodePrefix::SIM107 => SuffixLength::Three,
             CheckCodePrefix::SIM11 => SuffixLength::Two,
+            CheckCodePrefix::SIM110 => SuffixLength::Three,
+            CheckCodePrefix::SIM111 => SuffixLength::Three,
             CheckCodePrefix::SIM117 => SuffixLength::Three,
             CheckCodePrefix::SIM118 => SuffixLength::Three,
             CheckCodePrefix::SIM2 => SuffixLength::One,