Avoid raising PT012 for simple with statements (#6081)

crates/ruff/resources/test/fixtures/flake8_pytest_style/PT012.py

@@ -11,6 +11,10 @@ async def test_ok_trivial_with():
         with context_manager_under_test():
             pass
 
+    with pytest.raises(ValueError):
+        with context_manager_under_test():
+            raise ValueError
+
     with pytest.raises(AttributeError):
         async with context_manager_under_test():
             pass
@@ -47,13 +51,10 @@ async def test_error_complex_statement():
         while True:
             [].size
 
-    with pytest.raises(AttributeError):
-        with context_manager_under_test():
-            [].size
-
     with pytest.raises(AttributeError):
         async with context_manager_under_test():
-            [].size
+            if True:
+                raise Exception
 
 
 def test_error_try():

crates/ruff/src/rules/flake8_pytest_style/rules/raises.rs

@@ -1,4 +1,4 @@
-use ruff_python_ast::helpers::find_keyword;
+use ruff_python_ast::helpers::{find_keyword, is_compound_statement};
 use rustpython_parser::ast::{self, Expr, Keyword, Ranged, Stmt, WithItem};
 
 use ruff_diagnostics::{Diagnostic, Violation};
@@ -84,12 +84,10 @@ fn is_pytest_raises(func: &Expr, semantic: &SemanticModel) -> bool {
 }
 
 const fn is_non_trivial_with_body(body: &[Stmt]) -> bool {
-    if body.len() > 1 {
-        true
-    } else if let Some(first_body_stmt) = body.first() {
-        !first_body_stmt.is_pass_stmt()
+    if let [stmt] = body {
+        is_compound_statement(stmt)
     } else {
-        false
+        true
     }
 }
 
@@ -124,37 +122,24 @@ pub(crate) fn complex_raises(
     items: &[WithItem],
     body: &[Stmt],
 ) {
-    let mut is_too_complex = false;
-
     let raises_called = items.iter().any(|item| match &item.context_expr {
         Expr::Call(ast::ExprCall { func, .. }) => is_pytest_raises(func, checker.semantic()),
         _ => false,
     });
 
     // Check body for `pytest.raises` context manager
     if raises_called {
-        if body.len() > 1 {
-            is_too_complex = true;
-        } else if let Some(first_stmt) = body.first() {
-            match first_stmt {
+        let is_too_complex = if let [stmt] = body {
+            match stmt {
                 Stmt::With(ast::StmtWith { body, .. })
                 | Stmt::AsyncWith(ast::StmtAsyncWith { body, .. }) => {
-                    if is_non_trivial_with_body(body) {
-                        is_too_complex = true;
-                    }
-                }
-                Stmt::If(_)
-                | Stmt::For(_)
-                | Stmt::Match(_)
-                | Stmt::AsyncFor(_)
-                | Stmt::While(_)
-                | Stmt::Try(_)
-                | Stmt::TryStar(_) => {
-                    is_too_complex = true;
+                    is_non_trivial_with_body(body)
                 }
-                _ => {}
+                stmt => is_compound_statement(stmt),
             }
-        }
+        } else {
+            true
+        };
 
         if is_too_complex {
             checker.diagnostics.push(Diagnostic::new(

crates/ruff/src/rules/flake8_pytest_style/snapshots/ruff__rules__flake8_pytest_style__tests__PT012.snap

@@ -1,35 +1,22 @@
 ---
 source: crates/ruff/src/rules/flake8_pytest_style/mod.rs
 ---
-PT012.py:28:5: PT012 `pytest.raises()` block should contain a single simple statement
+PT012.py:32:5: PT012 `pytest.raises()` block should contain a single simple statement
    |
-27 |   def test_error_multiple_statements():
-28 |       with pytest.raises(AttributeError):
+31 |   def test_error_multiple_statements():
+32 |       with pytest.raises(AttributeError):
    |  _____^
-29 | |         len([])
-30 | |         [].size
+33 | |         len([])
+34 | |         [].size
    | |_______________^ PT012
    |
 
-PT012.py:34:5: PT012 `pytest.raises()` block should contain a single simple statement
-   |
-33 |   async def test_error_complex_statement():
-34 |       with pytest.raises(AttributeError):
-   |  _____^
-35 | |         if True:
-36 | |             [].size
-   | |___________________^ PT012
-37 |   
-38 |       with pytest.raises(AttributeError):
-   |
-
 PT012.py:38:5: PT012 `pytest.raises()` block should contain a single simple statement
    |
-36 |               [].size
-37 |   
+37 |   async def test_error_complex_statement():
 38 |       with pytest.raises(AttributeError):
    |  _____^
-39 | |         for i in []:
+39 | |         if True:
 40 | |             [].size
    | |___________________^ PT012
 41 |   
@@ -42,7 +29,7 @@ PT012.py:42:5: PT012 `pytest.raises()` block should contain a single simple stat
 41 |   
 42 |       with pytest.raises(AttributeError):
    |  _____^
-43 | |         async for i in []:
+43 | |         for i in []:
 44 | |             [].size
    | |___________________^ PT012
 45 |   
@@ -55,7 +42,7 @@ PT012.py:46:5: PT012 `pytest.raises()` block should contain a single simple stat
 45 |   
 46 |       with pytest.raises(AttributeError):
    |  _____^
-47 | |         while True:
+47 | |         async for i in []:
 48 | |             [].size
    | |___________________^ PT012
 49 |   
@@ -68,7 +55,7 @@ PT012.py:50:5: PT012 `pytest.raises()` block should contain a single simple stat
 49 |   
 50 |       with pytest.raises(AttributeError):
    |  _____^
-51 | |         with context_manager_under_test():
+51 | |         while True:
 52 | |             [].size
    | |___________________^ PT012
 53 |   
@@ -82,19 +69,20 @@ PT012.py:54:5: PT012 `pytest.raises()` block should contain a single simple stat
 54 |       with pytest.raises(AttributeError):
    |  _____^
 55 | |         async with context_manager_under_test():
-56 | |             [].size
-   | |___________________^ PT012
+56 | |             if True:
+57 | |                 raise Exception
+   | |_______________________________^ PT012
    |
 
-PT012.py:60:5: PT012 `pytest.raises()` block should contain a single simple statement
+PT012.py:61:5: PT012 `pytest.raises()` block should contain a single simple statement
    |
-59 |   def test_error_try():
-60 |       with pytest.raises(AttributeError):
+60 |   def test_error_try():
+61 |       with pytest.raises(AttributeError):
    |  _____^
-61 | |         try:
-62 | |             [].size
-63 | |         except:
-64 | |             raise
+62 | |         try:
+63 | |             [].size
+64 | |         except:
+65 | |             raise
    | |_________________^ PT012
    |
 

crates/ruff_python_ast/src/helpers.rs

@@ -18,6 +18,25 @@ use crate::call_path::CallPath;
 use crate::source_code::{Indexer, Locator};
 use crate::statement_visitor::{walk_body, walk_stmt, StatementVisitor};
 
+/// Return `true` if the `Stmt` is a compound statement (as opposed to a simple statement).
+pub const fn is_compound_statement(stmt: &Stmt) -> bool {
+    matches!(
+        stmt,
+        Stmt::FunctionDef(_)
+            | Stmt::AsyncFunctionDef(_)
+            | Stmt::ClassDef(_)
+            | Stmt::While(_)
+            | Stmt::For(_)
+            | Stmt::AsyncFor(_)
+            | Stmt::Match(_)
+            | Stmt::With(_)
+            | Stmt::AsyncWith(_)
+            | Stmt::If(_)
+            | Stmt::Try(_)
+            | Stmt::TryStar(_)
+    )
+}
+
 fn is_iterable_initializer<F>(id: &str, is_builtin: F) -> bool
 where
     F: Fn(&str) -> bool,

crates/ruff_python_formatter/src/statement/suite.rs

@@ -3,6 +3,7 @@ use rustpython_parser::ast::{Ranged, Stmt, Suite};
 use ruff_formatter::{
     format_args, write, FormatOwnedWithRule, FormatRefWithRule, FormatRuleWithOptions,
 };
+use ruff_python_ast::helpers::is_compound_statement;
 use ruff_python_trivia::lines_before;
 
 use crate::context::NodeLevel;
@@ -142,24 +143,6 @@ const fn is_class_or_function_definition(stmt: &Stmt) -> bool {
     )
 }
 
-const fn is_compound_statement(stmt: &Stmt) -> bool {
-    matches!(
-        stmt,
-        Stmt::FunctionDef(_)
-            | Stmt::AsyncFunctionDef(_)
-            | Stmt::ClassDef(_)
-            | Stmt::While(_)
-            | Stmt::For(_)
-            | Stmt::AsyncFor(_)
-            | Stmt::Match(_)
-            | Stmt::With(_)
-            | Stmt::AsyncWith(_)
-            | Stmt::If(_)
-            | Stmt::Try(_)
-            | Stmt::TryStar(_)
-    )
-}
-
 impl FormatRuleWithOptions<Suite, PyFormatContext<'_>> for FormatSuite {
     type Options = SuiteLevel;