diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_return/RET504.py b/crates/ruff_linter/resources/test/fixtures/flake8_return/RET504.py index 9899fa15ed6bc..8480aafaa1c01 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_return/RET504.py +++ b/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 diff --git a/crates/ruff_linter/src/rules/flake8_return/rules/function.rs b/crates/ruff_linter/src/rules/flake8_return/rules/function.rs index 7e016ecba70c7..4604663b3ece9 100644 --- a/crates/ruff_linter/src/rules/flake8_return/rules/function.rs +++ b/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); } diff --git a/crates/ruff_linter/src/rules/flake8_return/snapshots/ruff_linter__rules__flake8_return__tests__RET504_RET504.py.snap b/crates/ruff_linter/src/rules/flake8_return/snapshots/ruff_linter__rules__flake8_return__tests__RET504_RET504.py.snap index 9427cd7a79d66..3656e71afbd2d 100644 --- a/crates/ruff_linter/src/rules/flake8_return/snapshots/ruff_linter__rules__flake8_return__tests__RET504_RET504.py.snap +++ b/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(): diff --git a/crates/ruff_linter/src/rules/flake8_return/visitor.rs b/crates/ruff_linter/src/rules/flake8_return/visitor.rs index 775c3356e5f72..653364ff1ab78 100644 --- a/crates/ruff_linter/src/rules/flake8_return/visitor.rs +++ b/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<'a> { +pub(super) struct ReturnVisitor<'semantic, 'data> { + /// 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<'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<'a> Visitor<'a> for ReturnVisitor<'a> { +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, .. }) => { - if let Some(stmt_assign) = body.last().and_then(Stmt::as_assign_stmt) { - self.stack - .assignment_return - .push((stmt_assign, stmt_return, stmt)); + Stmt::With(with) => { + if let Some(stmt_assign) = + with.body.last().and_then(Stmt::as_assign_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 + }) +}