From c13f86aaf910b3876fb06b6668296bd3766886f7 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Mon, 14 Aug 2023 16:34:11 -0400
Subject: [PATCH] Use a visitor

---
 .../test/fixtures/perflint/PERF203.py         |  9 ++++
 .../perflint/rules/try_except_in_loop.rs      | 42 ++++++++++++++-----
 2 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/crates/ruff/resources/test/fixtures/perflint/PERF203.py b/crates/ruff/resources/test/fixtures/perflint/PERF203.py
index 5c6a79952e681..8a82c3c249ee2 100644
--- a/crates/ruff/resources/test/fixtures/perflint/PERF203.py
+++ b/crates/ruff/resources/test/fixtures/perflint/PERF203.py
@@ -38,3 +38,12 @@
     except:
         print("error")
 
+
+# OK - no other way to write this
+for i in range(10):
+    try:
+        print(f"{i}")
+        if i > 0:
+            break
+    except:
+        print("error")
diff --git a/crates/ruff/src/rules/perflint/rules/try_except_in_loop.rs b/crates/ruff/src/rules/perflint/rules/try_except_in_loop.rs
index 1d42e593d4758..84e14e5386a6a 100644
--- a/crates/ruff/src/rules/perflint/rules/try_except_in_loop.rs
+++ b/crates/ruff/src/rules/perflint/rules/try_except_in_loop.rs
@@ -1,7 +1,7 @@
-use ruff_python_ast::{self as ast, Ranged, Stmt};
-
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::statement_visitor::{walk_stmt, StatementVisitor};
+use ruff_python_ast::{self as ast, Ranged, Stmt};
 
 use crate::checkers::ast::Checker;
 use crate::settings::types::PythonVersion;
@@ -35,6 +35,7 @@ use crate::settings::types::PythonVersion;
 ///         int_numbers.append(int(num))
 ///     except ValueError as e:
 ///         print(f"Couldn't convert to integer: {e}")
+///         break
 /// ```
 ///
 /// Use instead:
@@ -67,7 +68,7 @@ pub(crate) fn try_except_in_loop(checker: &mut Checker, body: &[Stmt]) {
         return;
     }
 
-    let [Stmt::Try(ast::StmtTry { handlers, body: try_body, .. })] = body else {
+    let [Stmt::Try(ast::StmtTry { handlers, body, .. })] = body else {
         return;
     };
 
@@ -75,16 +76,37 @@ pub(crate) fn try_except_in_loop(checker: &mut Checker, body: &[Stmt]) {
         return;
     };
 
-    // Parse all nodes underneath the `try` - if we see any loop control flow statements, we don't
-    // want to throw.
-    for stmt in try_body {
-        match stmt {
-            Stmt::Continue(_) | Stmt::Break(_) => return,
-            _ => (),
-        }
+    // Avoid flagging `try`-`except` blocks that contain `break` or `continue`,
+    // which rely on the exception handling mechanism.
+    if has_break_or_continue(body) {
+        return;
     }
 
     checker
         .diagnostics
         .push(Diagnostic::new(TryExceptInLoop, handler.range()));
 }
+
+/// Returns `true` if a `break` or `continue` statement is present in `body`.
+fn has_break_or_continue(body: &[Stmt]) -> bool {
+    let mut visitor = LoopControlFlowVisitor::default();
+    visitor.visit_body(body);
+    visitor.has_break_or_continue
+}
+
+#[derive(Debug, Default)]
+struct LoopControlFlowVisitor {
+    has_break_or_continue: bool,
+}
+
+impl StatementVisitor<'_> for LoopControlFlowVisitor {
+    fn visit_stmt(&mut self, stmt: &Stmt) {
+        match stmt {
+            Stmt::Break(_) | Stmt::Continue(_) => self.has_break_or_continue = true,
+            Stmt::FunctionDef(_) | Stmt::ClassDef(_) => {
+                // Don't recurse.
+            }
+            _ => walk_stmt(self, stmt),
+        }
+    }
+}