Use a visitor

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")

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,24 +68,40 @@ 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;
     };
 
     let Some(handler) = handlers.first() else {
         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.
+    let mut visitor = LoopControlFlowVisitor::default();
+    visitor.visit_body(body);
+    if visitor.has_break_or_continue {
+        return;
     }
 
     checker
         .diagnostics
         .push(Diagnostic::new(TryExceptInLoop, handler.range()));
 }
+
+#[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),
+        }
+    }
+}