From c6c15d5cf9512ca8e2f39d38ca7884c20de3fa0b Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Mon, 13 Feb 2023 21:51:12 -0500
Subject: [PATCH] Avoid unnecessary-else violations in `elif` branches (#2881)

Long-time source of confusion -- two reports over 1800 issues apart.

Closes #1035.

Closes #2879.
---
 .../test/fixtures/flake8_return/RET508.py     | 21 ++++++++++
 crates/ruff/src/rules/flake8_return/rules.rs  |  8 +---
 .../ruff/src/rules/flake8_return/visitor.rs   | 42 ++++++++++++++++---
 3 files changed, 58 insertions(+), 13 deletions(-)

diff --git a/crates/ruff/resources/test/fixtures/flake8_return/RET508.py b/crates/ruff/resources/test/fixtures/flake8_return/RET508.py
index 40f053bfa5..156b3a0250 100644
--- a/crates/ruff/resources/test/fixtures/flake8_return/RET508.py
+++ b/crates/ruff/resources/test/fixtures/flake8_return/RET508.py
@@ -136,3 +136,24 @@ def nested2(x, y, z):
                 break
         else:
             a = z
+
+
+def elif1(x, y, w, z):
+    for i in x:
+        if i > y:
+            a = z
+        elif i > w:
+            break
+        else:
+            a = z
+
+
+def elif2(x, y, w, z):
+    for i in x:
+        if i > y:
+            a = z
+        else:
+            if i > w:
+                break
+            else:
+                a = z
diff --git a/crates/ruff/src/rules/flake8_return/rules.rs b/crates/ruff/src/rules/flake8_return/rules.rs
index 04e71d495e..6801be192d 100644
--- a/crates/ruff/src/rules/flake8_return/rules.rs
+++ b/crates/ruff/src/rules/flake8_return/rules.rs
@@ -423,13 +423,7 @@ fn superfluous_elif(checker: &mut Checker, stack: &Stack) -> bool {
 
 /// RET505, RET506, RET507, RET508
 fn superfluous_else(checker: &mut Checker, stack: &Stack) -> bool {
-    for stmt in &stack.ifs {
-        let StmtKind::If { orelse, .. } = &stmt.node else {
-            continue;
-        };
-        if orelse.is_empty() {
-            continue;
-        }
+    for stmt in &stack.elses {
         if superfluous_else_node(checker, stmt, Branch::Else) {
             return true;
         }
diff --git a/crates/ruff/src/rules/flake8_return/visitor.rs b/crates/ruff/src/rules/flake8_return/visitor.rs
index 002244025d..74768eb91e 100644
--- a/crates/ruff/src/rules/flake8_return/visitor.rs
+++ b/crates/ruff/src/rules/flake8_return/visitor.rs
@@ -8,7 +8,7 @@ use crate::ast::visitor::Visitor;
 pub struct Stack<'a> {
     pub returns: Vec<(&'a Stmt, Option<&'a Expr>)>,
     pub yields: Vec<&'a Expr>,
-    pub ifs: Vec<&'a Stmt>,
+    pub elses: Vec<&'a Stmt>,
     pub elifs: Vec<&'a Stmt>,
     pub refs: FxHashMap<&'a str, Vec<Location>>,
     pub non_locals: FxHashSet<&'a str>,
@@ -20,6 +20,7 @@ pub struct Stack<'a> {
 #[derive(Default)]
 pub struct ReturnVisitor<'a> {
     pub stack: Stack<'a>,
+    parents: Vec<&'a Stmt>,
 }
 
 impl<'a> ReturnVisitor<'a> {
@@ -72,16 +73,37 @@ impl<'a> Visitor<'a> for ReturnVisitor<'a> {
                 self.stack
                     .returns
                     .push((stmt, value.as_ref().map(|expr| &**expr)));
+
+                self.parents.push(stmt);
                 visitor::walk_stmt(self, stmt);
+                self.parents.pop();
             }
             StmtKind::If { orelse, .. } => {
-                if orelse.len() == 1 && matches!(orelse.first().unwrap().node, StmtKind::If { .. })
-                {
-                    self.stack.elifs.push(stmt);
-                } else {
-                    self.stack.ifs.push(stmt);
+                let is_elif_arm = self.parents.iter().any(|parent| {
+                    if let StmtKind::If { orelse, .. } = &parent.node {
+                        orelse.len() == 1 && &orelse[0] == stmt
+                    } else {
+                        false
+                    }
+                });
+
+                if !is_elif_arm {
+                    let has_elif = orelse.len() == 1
+                        && matches!(orelse.first().unwrap().node, StmtKind::If { .. });
+                    let has_else = !orelse.is_empty();
+
+                    if has_elif {
+                        // `stmt` is an `if` block followed by an `elif` clause.
+                        self.stack.elifs.push(stmt);
+                    } else if has_else {
+                        // `stmt` is an `if` block followed by an `else` clause.
+                        self.stack.elses.push(stmt);
+                    }
                 }
+
+                self.parents.push(stmt);
                 visitor::walk_stmt(self, stmt);
+                self.parents.pop();
             }
             StmtKind::Assign { targets, value, .. } => {
                 if let ExprKind::Name { id, .. } = &value.node {
@@ -109,16 +131,24 @@ impl<'a> Visitor<'a> for ReturnVisitor<'a> {
                 self.stack
                     .loops
                     .push((stmt.location, stmt.end_location.unwrap()));
+
+                self.parents.push(stmt);
                 visitor::walk_stmt(self, stmt);
+                self.parents.pop();
             }
             StmtKind::Try { .. } => {
                 self.stack
                     .tries
                     .push((stmt.location, stmt.end_location.unwrap()));
+
+                self.parents.push(stmt);
                 visitor::walk_stmt(self, stmt);
+                self.parents.pop();
             }
             _ => {
+                self.parents.push(stmt);
                 visitor::walk_stmt(self, stmt);
+                self.parents.pop();
             }
         }
     }