From 74cdfdb8e689b5ad7f3967387228f0a89d305f7d Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Fri, 28 Jul 2023 00:55:55 -0400
Subject: [PATCH] Skip `PERF203` violations for multi-statement loops (#6145)

Closes https://github.com/astral-sh/ruff/issues/5858.
---
 .../test/fixtures/perflint/PERF203.py         | 15 ++++-------
 .../perflint/rules/try_except_in_loop.rs      | 21 ++++++++-------
 ...__perflint__tests__PERF203_PERF203.py.snap | 26 +++++--------------
 3 files changed, 23 insertions(+), 39 deletions(-)

diff --git a/crates/ruff/resources/test/fixtures/perflint/PERF203.py b/crates/ruff/resources/test/fixtures/perflint/PERF203.py
index ec3ca5feeefaf..14b3fa38a3800 100644
--- a/crates/ruff/resources/test/fixtures/perflint/PERF203.py
+++ b/crates/ruff/resources/test/fixtures/perflint/PERF203.py
@@ -1,28 +1,23 @@
+# PERF203
 for i in range(10):
-    try:  # PERF203
+    try:
         print(f"{i}")
     except:
         print("error")
 
+# OK
 try:
     for i in range(10):
         print(f"{i}")
 except:
     print("error")
 
+# OK
 i = 0
-while i < 10:  # PERF203
+while i < 10:
     try:
         print(f"{i}")
     except:
         print("error")
 
     i += 1
-
-try:
-    i = 0
-    while i < 10:
-        print(f"{i}")
-        i += 1
-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 fe57321d8ca3a..ffa8a428faa3e 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
@@ -67,14 +67,15 @@ pub(crate) fn try_except_in_loop(checker: &mut Checker, body: &[Stmt]) {
         return;
     }
 
-    checker.diagnostics.extend(body.iter().filter_map(|stmt| {
-        if let Stmt::Try(ast::StmtTry { handlers, .. }) = stmt {
-            handlers
-                .iter()
-                .next()
-                .map(|handler| Diagnostic::new(TryExceptInLoop, handler.range()))
-        } else {
-            None
-        }
-    }));
+    let [Stmt::Try(ast::StmtTry { handlers, .. })] = body else {
+        return;
+    };
+
+    let Some(handler) = handlers.first() else {
+        return;
+    };
+
+    checker
+        .diagnostics
+        .push(Diagnostic::new(TryExceptInLoop, handler.range()));
 }
diff --git a/crates/ruff/src/rules/perflint/snapshots/ruff__rules__perflint__tests__PERF203_PERF203.py.snap b/crates/ruff/src/rules/perflint/snapshots/ruff__rules__perflint__tests__PERF203_PERF203.py.snap
index 08a287819e7ab..d74621c223552 100644
--- a/crates/ruff/src/rules/perflint/snapshots/ruff__rules__perflint__tests__PERF203_PERF203.py.snap
+++ b/crates/ruff/src/rules/perflint/snapshots/ruff__rules__perflint__tests__PERF203_PERF203.py.snap
@@ -1,28 +1,16 @@
 ---
 source: crates/ruff/src/rules/perflint/mod.rs
 ---
-PERF203.py:4:5: PERF203 `try`-`except` within a loop incurs performance overhead
+PERF203.py:5:5: PERF203 `try`-`except` within a loop incurs performance overhead
   |
-2 |       try:  # PERF203
-3 |           print(f"{i}")
-4 |       except:
+3 |       try:
+4 |           print(f"{i}")
+5 |       except:
   |  _____^
-5 | |         print("error")
+6 | |         print("error")
   | |______________________^ PERF203
-6 |   
-7 |   try:
+7 |   
+8 |   # OK
   |
 
-PERF203.py:17:5: PERF203 `try`-`except` within a loop incurs performance overhead
-   |
-15 |       try:
-16 |           print(f"{i}")
-17 |       except:
-   |  _____^
-18 | |         print("error")
-   | |______________________^ PERF203
-19 |   
-20 |       i += 1
-   |
-