diff --git a/README.md b/README.md index f7521e6d2c..ff2957a3dc 100644 --- a/README.md +++ b/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 ...` | 🛠 | diff --git a/resources/test/fixtures/flake8_simplify/SIM110.py b/resources/test/fixtures/flake8_simplify/SIM110.py new file mode 100644 index 0000000000..7f7daa05b0 --- /dev/null +++ b/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" diff --git a/resources/test/fixtures/flake8_simplify/SIM111.py b/resources/test/fixtures/flake8_simplify/SIM111.py new file mode 100644 index 0000000000..eb92c7fdce --- /dev/null +++ b/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" diff --git a/src/checkers/ast.rs b/src/checkers/ast.rs index baf2573dfa..4eae376a12 100644 --- a/src/checkers/ast.rs +++ b/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); } } diff --git a/src/flake8_simplify/mod.rs b/src/flake8_simplify/mod.rs index 4979308043..91f8bc0f83 100644 --- a/src/flake8_simplify/mod.rs +++ b/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")] diff --git a/src/flake8_simplify/plugins/ast_for.rs b/src/flake8_simplify/plugins/ast_for.rs new file mode 100644 index 0000000000..1787984fc9 --- /dev/null +++ b/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> { + 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 { + 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); + } + } + } + } +} diff --git a/src/flake8_simplify/plugins/mod.rs b/src/flake8_simplify/plugins/mod.rs index 9634abd8bc..8f7c337cf3 100644 --- a/src/flake8_simplify/plugins/mod.rs +++ b/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; 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 new file mode 100644 index 0000000000..ac46076449 --- /dev/null +++ b/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: ~ + 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 new file mode 100644 index 0000000000..05cbdbd38a --- /dev/null +++ b/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: ~ + diff --git a/src/registry.rs b/src/registry.rs index 656963ed18..c651c6186d 100644 --- a/src/registry.rs +++ b/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")) diff --git a/src/registry_gen.rs b/src/registry_gen.rs index 0117bea98e..77f1803884 100644 --- a/src/registry_gen.rs +++ b/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,