[flake8-return] Consider exception suppress for unnecessary assignment (#9673)

## Summary

This review contains a fix for
[RET504](https://docs.astral.sh/ruff/rules/unnecessary-assign/)
(unnecessary-assign)

The problem is that Ruff suggests combining a return statement inside
contextlib.suppress. Even though it is an unsafe fix it might lead to an
invalid code that is not equivalent to the original one.

See: https://github.com/astral-sh/ruff/issues/5909

## Test Plan

```bash
cargo test
```

crates/ruff_linter/resources/test/fixtures/flake8_return/RET504.py

@@ -363,3 +363,46 @@ def foo():
 def mavko_debari(P_kbar):
     D=0.4853881 + 3.6006116*P - 0.0117368*(P-1.3822)**2
     return D
+
+
+# contextlib suppress in with statement
+import contextlib
+
+
+def foo():
+    x = 2
+    with contextlib.suppress(Exception):
+        x = x + 1
+    return x
+
+
+def foo(data):
+    with open("in.txt") as file_out, contextlib.suppress(IOError):
+        file_out.write(data)
+        data = 10
+    return data
+
+
+def foo(data):
+    with open("in.txt") as file_out:
+        file_out.write(data)
+        with contextlib.suppress(IOError):
+            data = 10
+    return data
+
+
+def foo():
+    y = 1
+    x = 2
+    with contextlib.suppress(Exception):
+        x = 1
+    y = y + 2
+    return y  # RET504
+
+
+def foo():
+    y = 1
+    if y > 0:
+        with contextlib.suppress(Exception):
+            y = 2
+        return y

crates/ruff_linter/src/rules/flake8_return/rules/function.rs

@@ -737,7 +737,7 @@ pub(crate) fn function(checker: &mut Checker, body: &[Stmt], returns: Option<&Ex
 
     // Traverse the function body, to collect the stack.
     let stack = {
-        let mut visitor = ReturnVisitor::default();
+        let mut visitor = ReturnVisitor::new(checker.semantic());
         for stmt in body {
             visitor.visit_stmt(stmt);
         }

crates/ruff_linter/src/rules/flake8_return/snapshots/ruff_linter__rules__flake8_return__tests__RET504_RET504.py.snap

@@ -217,5 +217,28 @@ RET504.py:365:12: RET504 [*] Unnecessary assignment to `D` before `return` state
 364     |-    D=0.4853881 + 3.6006116*P - 0.0117368*(P-1.3822)**2
 365     |-    return D
     364 |+    return 0.4853881 + 3.6006116*P - 0.0117368*(P-1.3822)**2
+366 365 | 
+367 366 | 
+368 367 | # contextlib suppress in with statement
+
+RET504.py:400:12: RET504 [*] Unnecessary assignment to `y` before `return` statement
+    |
+398 |         x = 1
+399 |     y = y + 2
+400 |     return y  # RET504
+    |            ^ RET504
+    |
+    = help: Remove unnecessary assignment
+
+ℹ Unsafe fix
+396 396 |     x = 2
+397 397 |     with contextlib.suppress(Exception):
+398 398 |         x = 1
+399     |-    y = y + 2
+400     |-    return y  # RET504
+    399 |+    return y + 2
+401 400 | 
+402 401 | 
+403 402 | def foo():
 
 

crates/ruff_linter/src/rules/flake8_return/visitor.rs

@@ -3,34 +3,48 @@ use rustc_hash::FxHashSet;
 
 use ruff_python_ast::visitor;
 use ruff_python_ast::visitor::Visitor;
+use ruff_python_semantic::SemanticModel;
 
 #[derive(Default)]
-pub(super) struct Stack<'a> {
+pub(super) struct Stack<'data> {
     /// The `return` statements in the current function.
-    pub(super) returns: Vec<&'a ast::StmtReturn>,
+    pub(super) returns: Vec<&'data ast::StmtReturn>,
     /// The `elif` or `else` statements in the current function.
-    pub(super) elifs_elses: Vec<(&'a [Stmt], &'a ElifElseClause)>,
+    pub(super) elifs_elses: Vec<(&'data [Stmt], &'data ElifElseClause)>,
     /// The non-local variables in the current function.
-    pub(super) non_locals: FxHashSet<&'a str>,
+    pub(super) non_locals: FxHashSet<&'data str>,
     /// Whether the current function is a generator.
     pub(super) is_generator: bool,
     /// The `assignment`-to-`return` statement pairs in the current function.
     /// TODO(charlie): Remove the extra [`Stmt`] here, which is necessary to support statement
     /// removal for the `return` statement.
-    pub(super) assignment_return: Vec<(&'a ast::StmtAssign, &'a ast::StmtReturn, &'a Stmt)>,
+    pub(super) assignment_return:
+        Vec<(&'data ast::StmtAssign, &'data ast::StmtReturn, &'data Stmt)>,
 }
 
-#[derive(Default)]
+pub(super) struct ReturnVisitor<'semantic, 'data> {
-pub(super) struct ReturnVisitor<'a> {
+    /// The semantic model of the current file.
+    semantic: &'semantic SemanticModel<'data>,
     /// The current stack of nodes.
-    pub(super) stack: Stack<'a>,
+    pub(super) stack: Stack<'data>,
     /// The preceding sibling of the current node.
-    sibling: Option<&'a Stmt>,
+    sibling: Option<&'data Stmt>,
     /// The parent nodes of the current node.
-    parents: Vec<&'a Stmt>,
+    parents: Vec<&'data Stmt>,
 }
 
-impl<'a> Visitor<'a> for ReturnVisitor<'a> {
+impl<'semantic, 'data> ReturnVisitor<'semantic, 'data> {
+    pub(super) fn new(semantic: &'semantic SemanticModel<'data>) -> Self {
+        Self {
+            semantic,
+            stack: Stack::default(),
+            sibling: None,
+            parents: Vec::new(),
+        }
+    }
+}
+
+impl<'semantic, 'a> Visitor<'a> for ReturnVisitor<'semantic, 'a> {
     fn visit_stmt(&mut self, stmt: &'a Stmt) {
         match stmt {
             Stmt::ClassDef(ast::StmtClassDef { decorator_list, .. }) => {
@@ -95,11 +109,17 @@ impl<'a> Visitor<'a> for ReturnVisitor<'a> {
                         //         x = f.read()
                         //     return x
                         // ```
-                        Stmt::With(ast::StmtWith { body, .. }) => {
+                        Stmt::With(with) => {
-                            if let Some(stmt_assign) = body.last().and_then(Stmt::as_assign_stmt) {
+                            if let Some(stmt_assign) =
-                                self.stack
+                                with.body.last().and_then(Stmt::as_assign_stmt)
-                                    .assignment_return
+                            {
-                                    .push((stmt_assign, stmt_return, stmt));
+                                if !has_conditional_body(with, self.semantic) {
+                                    self.stack.assignment_return.push((
+                                        stmt_assign,
+                                        stmt_return,
+                                        stmt,
+                                    ));
+                                }
                             }
                         }
                         _ => {}
@@ -142,3 +162,47 @@ impl<'a> Visitor<'a> for ReturnVisitor<'a> {
         self.sibling = sibling;
     }
 }
+
+/// RET504
+/// If the last statement is a `return` statement, and the second-to-last statement is a
+/// `with` statement that suppresses an exception, then we should not analyze the `return`
+/// statement for unnecessary assignments. Otherwise we will suggest removing the assignment
+/// and the `with` statement, which would change the behavior of the code.
+///
+/// Example:
+/// ```python
+/// def foo(data):
+///    with suppress(JSONDecoderError):
+///       data = data.decode()
+///   return data
+
+/// Returns `true` if the [`With`] statement is known to have a conditional body. In other words:
+/// if the [`With`] statement's body may or may not run.
+///
+/// For example, in the following, it's unsafe to inline the `return` into the `with`, since if
+/// `data.decode()` fails, the behavior of the program will differ. (As-is, the function will return
+/// the input `data`; if we inline the `return`, the function will return `None`.)
+///
+/// ```python
+/// def func(data):
+///     with suppress(JSONDecoderError):
+///         data = data.decode()
+///     return data
+/// ```
+fn has_conditional_body(with: &ast::StmtWith, semantic: &SemanticModel) -> bool {
+    with.items.iter().any(|item| {
+        let ast::WithItem {
+            context_expr: Expr::Call(ast::ExprCall { func, .. }),
+            ..
+        } = item
+        else {
+            return false;
+        };
+        if let Some(call_path) = semantic.resolve_call_path(func) {
+            if call_path.as_slice() == ["contextlib", "suppress"] {
+                return true;
+            }
+        }
+        false
+    })
+}