feat: let SIM103 return expressions without bool() wrapping

Issue astral-sh#2240
spaceone · Jan 31, 2023 · dd03891 · dd03891
1 parent 15d4774
commit dd03891
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 20 deletions.
diff --git a/resources/test/fixtures/flake8_simplify/SIM103.py b/resources/test/fixtures/flake8_simplify/SIM103.py
@@ -6,6 +6,14 @@ def f():
         return False
 
 
+def f():
+    # SIM103
+    if a == b:
+        return True
+    else:
+        return False
+
+
 def f():
     # SIM103
     if a:

diff --git a/src/rules/flake8_simplify/rules/ast_if.rs b/src/rules/flake8_simplify/rules/ast_if.rs
@@ -184,16 +184,26 @@ pub fn return_bool_condition_directly(checker: &mut Checker, stmt: &Stmt) {
         && matches!(else_return, Bool::False)
         && !has_comments(stmt, checker.locator)
     {
-        let return_stmt = create_stmt(StmtKind::Return {
-            value: Some(Box::new(create_expr(ExprKind::Call {
-                func: Box::new(create_expr(ExprKind::Name {
-                    id: "bool".to_string(),
-                    ctx: ExprContext::Load,
-                })),
-                args: vec![(**test).clone()],
-                keywords: vec![],
-            }))),
-        });
+        let return_stmt;
+        match test.node {
+            ExprKind::Compare { .. } => {
+                return_stmt = create_stmt(StmtKind::Return {
+                    value: Some(test.clone()),
+                })
+            }
+            _ => {
+                return_stmt = create_stmt(StmtKind::Return {
+                    value: Some(Box::new(create_expr(ExprKind::Call {
+                        func: Box::new(create_expr(ExprKind::Name {
+                            id: "bool".to_string(),
+                            ctx: ExprContext::Load,
+                        })),
+                        args: vec![(**test).clone()],
+                        keywords: vec![],
+                    }))),
+                })
+            }
+        }
         diagnostic.amend(Fix::replacement(
             unparse_stmt(&return_stmt, checker.stylist),
             stmt.location,

diff --git a/...ules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM103_SIM103.py.snap b/...ules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM103_SIM103.py.snap
@@ -21,52 +21,71 @@ expression: diagnostics
       row: 6
       column: 20
   parent: ~
+- kind:
+    ReturnBoolConditionDirectly:
+      cond: a == b
+  location:
+    row: 11
+    column: 4
+  end_location:
+    row: 14
+    column: 20
+  fix:
+    content:
+      - return a == b
+    location:
+      row: 11
+      column: 4
+    end_location:
+      row: 14
+      column: 20
+  parent: ~
 - kind:
     ReturnBoolConditionDirectly:
       cond: b
   location:
-    row: 13
+    row: 21
     column: 4
   end_location:
-    row: 16
+    row: 24
     column: 20
   fix:
     content:
       - return bool(b)
     location:
-      row: 13
+      row: 21
       column: 4
     end_location:
-      row: 16
+      row: 24
       column: 20
   parent: ~
 - kind:
     ReturnBoolConditionDirectly:
       cond: b
   location:
-    row: 24
+    row: 32
     column: 8
   end_location:
-    row: 27
+    row: 35
     column: 24
   fix:
     content:
       - return bool(b)
     location:
-      row: 24
+      row: 32
       column: 8
     end_location:
-      row: 27
+      row: 35
       column: 24
   parent: ~
 - kind:
     ReturnBoolConditionDirectly:
       cond: a
   location:
-    row: 49
+    row: 57
     column: 4
   end_location:
-    row: 52
+    row: 60
     column: 19
   fix: ~
   parent: ~