From 84d1df08bebf0510e02acef50bdfc1972ad7e7ed Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Tue, 17 Jan 2023 22:03:17 -0500
Subject: [PATCH] Avoid broken autofix for `SIM103` with `elif` (#1944)

Also adjusts the generator to avoid the extra parentheses (and skips commented `if` statements).

Closes #1943.
---
 .../test/fixtures/flake8_simplify/SIM103.py   | 30 +++++++++++--
 src/rules/flake8_simplify/rules/ast_if.rs     |  6 ++-
 ...ke8_simplify__tests__SIM103_SIM103.py.snap | 42 +++++++++++++++++--
 ...les__pycodestyle__tests__E731_E731.py.snap |  6 +--
 src/source_code/generator.rs                  |  5 ++-
 5 files changed, 76 insertions(+), 13 deletions(-)

diff --git a/resources/test/fixtures/flake8_simplify/SIM103.py b/resources/test/fixtures/flake8_simplify/SIM103.py
index 7fa8d66c1e..5b72e0947d 100644
--- a/resources/test/fixtures/flake8_simplify/SIM103.py
+++ b/resources/test/fixtures/flake8_simplify/SIM103.py
@@ -1,12 +1,35 @@
 def f():
-    if a:  # SIM103
+    # SIM103
+    if a:
         return True
     else:
         return False
 
 
 def f():
-    if a:  # OK
+    # SIM103
+    if a:
+        return 1
+    elif b:
+        return True
+    else:
+        return False
+
+
+def f():
+    # SIM103
+    if a:
+        return 1
+    else:
+        if b:
+            return True
+        else:
+            return False
+
+
+def f():
+    # OK
+    if a:
         foo()
         return True
     else:
@@ -14,7 +37,8 @@ def f():
 
 
 def f():
-    if a:  # OK
+    # OK
+    if a:
         return "foo"
     else:
         return False
diff --git a/src/rules/flake8_simplify/rules/ast_if.rs b/src/rules/flake8_simplify/rules/ast_if.rs
index 62465314cb..b98ed7872e 100644
--- a/src/rules/flake8_simplify/rules/ast_if.rs
+++ b/src/rules/flake8_simplify/rules/ast_if.rs
@@ -96,14 +96,16 @@ pub fn return_bool_condition_directly(checker: &mut Checker, stmt: &Stmt) {
         violations::ReturnBoolConditionDirectly(condition),
         Range::from_located(stmt),
     );
-    if checker.patch(&RuleCode::SIM103) {
+    if checker.patch(&RuleCode::SIM103)
+        && !(has_comments(stmt, checker.locator) || has_comments(&orelse[0], checker.locator))
+    {
         let return_stmt = create_stmt(StmtKind::Return {
             value: Some(test.clone()),
         });
         diagnostic.amend(Fix::replacement(
             unparse_stmt(&return_stmt, checker.stylist),
             stmt.location,
-            stmt.end_location.unwrap(),
+            orelse[0].end_location.unwrap(),
         ));
     }
     checker.diagnostics.push(diagnostic);
diff --git a/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM103_SIM103.py.snap b/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM103_SIM103.py.snap
index 1fcbffcb08..e48c0866cb 100644
--- a/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM103_SIM103.py.snap
+++ b/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM103_SIM103.py.snap
@@ -5,18 +5,52 @@ expression: diagnostics
 - kind:
     ReturnBoolConditionDirectly: a
   location:
-    row: 2
+    row: 3
     column: 4
   end_location:
-    row: 5
+    row: 6
     column: 20
   fix:
     content: return a
     location:
-      row: 2
+      row: 3
       column: 4
     end_location:
-      row: 5
+      row: 6
       column: 20
   parent: ~
+- kind:
+    ReturnBoolConditionDirectly: b
+  location:
+    row: 13
+    column: 4
+  end_location:
+    row: 14
+    column: 19
+  fix:
+    content: return b
+    location:
+      row: 13
+      column: 4
+    end_location:
+      row: 16
+      column: 20
+  parent: ~
+- kind:
+    ReturnBoolConditionDirectly: b
+  location:
+    row: 24
+    column: 8
+  end_location:
+    row: 27
+    column: 24
+  fix:
+    content: return b
+    location:
+      row: 24
+      column: 8
+    end_location:
+      row: 27
+      column: 24
+  parent: ~
 
diff --git a/src/rules/pycodestyle/snapshots/ruff__rules__pycodestyle__tests__E731_E731.py.snap b/src/rules/pycodestyle/snapshots/ruff__rules__pycodestyle__tests__E731_E731.py.snap
index ee8611c78c..66c4f7a01c 100644
--- a/src/rules/pycodestyle/snapshots/ruff__rules__pycodestyle__tests__E731_E731.py.snap
+++ b/src/rules/pycodestyle/snapshots/ruff__rules__pycodestyle__tests__E731_E731.py.snap
@@ -11,7 +11,7 @@ expression: diagnostics
     row: 2
     column: 19
   fix:
-    content: "def f(x):\n    return (2 * x)"
+    content: "def f(x):\n    return 2 * x"
     location:
       row: 2
       column: 0
@@ -28,7 +28,7 @@ expression: diagnostics
     row: 4
     column: 19
   fix:
-    content: "def f(x):\n    return (2 * x)"
+    content: "def f(x):\n    return 2 * x"
     location:
       row: 4
       column: 0
@@ -45,7 +45,7 @@ expression: diagnostics
     row: 7
     column: 29
   fix:
-    content: "def this(y, z):\n        return (2 * x)"
+    content: "def this(y, z):\n        return 2 * x"
     location:
       row: 7
       column: 4
diff --git a/src/source_code/generator.rs b/src/source_code/generator.rs
index 560ec5fb19..6843e4c553 100644
--- a/src/source_code/generator.rs
+++ b/src/source_code/generator.rs
@@ -227,10 +227,13 @@ impl<'a> Generator<'a> {
                 }
             }
             StmtKind::Return { value } => {
+                // TODO(charlie): Revisit precedence. In particular, look at how Astor handles
+                // precedence.
+                // See: https://github.com/berkerpeksag/astor/blob/8342d6aa5dcdcf20f89a19057527510c245c7a2e/astor/code_gen.py#L86
                 statement!({
                     if let Some(expr) = value {
                         self.p("return ");
-                        self.unparse_expr(expr, precedence::ATOM);
+                        self.unparse_expr(expr, 1);
                     } else {
                         self.p("return");
                     }