Skip to content

Commit ad2cfa3

Browse files
mikeleppanemikeleppane
authored
[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: #5909 ## Test Plan ```bash cargo test ```
1 parent 0045032 commit ad2cfa3

4 files changed

Lines changed: 147 additions & 17 deletions

File tree

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

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -363,3 +363,46 @@ def foo():
363363
def mavko_debari(P_kbar):
364364
D=0.4853881 + 3.6006116*P - 0.0117368*(P-1.3822)**2
365365
return D
366+
367+
368+
# contextlib suppress in with statement
369+
import contextlib
370+
371+
372+
def foo():
373+
x = 2
374+
with contextlib.suppress(Exception):
375+
x = x + 1
376+
return x
377+
378+
379+
def foo(data):
380+
with open("in.txt") as file_out, contextlib.suppress(IOError):
381+
file_out.write(data)
382+
data = 10
383+
return data
384+
385+
386+
def foo(data):
387+
with open("in.txt") as file_out:
388+
file_out.write(data)
389+
with contextlib.suppress(IOError):
390+
data = 10
391+
return data
392+
393+
394+
def foo():
395+
y = 1
396+
x = 2
397+
with contextlib.suppress(Exception):
398+
x = 1
399+
y = y + 2
400+
return y # RET504
401+
402+
403+
def foo():
404+
y = 1
405+
if y > 0:
406+
with contextlib.suppress(Exception):
407+
y = 2
408+
return y

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -737,7 +737,7 @@ pub(crate) fn function(checker: &mut Checker, body: &[Stmt], returns: Option<&Ex
737737

738738
// Traverse the function body, to collect the stack.
739739
let stack = {
740-
let mut visitor = ReturnVisitor::default();
740+
let mut visitor = ReturnVisitor::new(checker.semantic());
741741
for stmt in body {
742742
visitor.visit_stmt(stmt);
743743
}

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,5 +217,28 @@ RET504.py:365:12: RET504 [*] Unnecessary assignment to `D` before `return` state
217217
364 |- D=0.4853881 + 3.6006116*P - 0.0117368*(P-1.3822)**2
218218
365 |- return D
219219
364 |+ return 0.4853881 + 3.6006116*P - 0.0117368*(P-1.3822)**2
220+
366 365 |
221+
367 366 |
222+
368 367 | # contextlib suppress in with statement
223+
224+
RET504.py:400:12: RET504 [*] Unnecessary assignment to `y` before `return` statement
225+
|
226+
398 | x = 1
227+
399 | y = y + 2
228+
400 | return y # RET504
229+
| ^ RET504
230+
|
231+
= help: Remove unnecessary assignment
232+
233+
ℹ Unsafe fix
234+
396 396 | x = 2
235+
397 397 | with contextlib.suppress(Exception):
236+
398 398 | x = 1
237+
399 |- y = y + 2
238+
400 |- return y # RET504
239+
399 |+ return y + 2
240+
401 400 |
241+
402 401 |
242+
403 402 | def foo():
220243

221244

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

Lines changed: 80 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -3,34 +3,48 @@ use rustc_hash::FxHashSet;
33

44
use ruff_python_ast::visitor;
55
use ruff_python_ast::visitor::Visitor;
6+
use ruff_python_semantic::SemanticModel;
67

78
#[derive(Default)]
8-
pub(super) struct Stack<'a> {
9+
pub(super) struct Stack<'data> {
910
/// The `return` statements in the current function.
10-
pub(super) returns: Vec<&'a ast::StmtReturn>,
11+
pub(super) returns: Vec<&'data ast::StmtReturn>,
1112
/// The `elif` or `else` statements in the current function.
12-
pub(super) elifs_elses: Vec<(&'a [Stmt], &'a ElifElseClause)>,
13+
pub(super) elifs_elses: Vec<(&'data [Stmt], &'data ElifElseClause)>,
1314
/// The non-local variables in the current function.
14-
pub(super) non_locals: FxHashSet<&'a str>,
15+
pub(super) non_locals: FxHashSet<&'data str>,
1516
/// Whether the current function is a generator.
1617
pub(super) is_generator: bool,
1718
/// The `assignment`-to-`return` statement pairs in the current function.
1819
/// TODO(charlie): Remove the extra [`Stmt`] here, which is necessary to support statement
1920
/// removal for the `return` statement.
20-
pub(super) assignment_return: Vec<(&'a ast::StmtAssign, &'a ast::StmtReturn, &'a Stmt)>,
21+
pub(super) assignment_return:
22+
Vec<(&'data ast::StmtAssign, &'data ast::StmtReturn, &'data Stmt)>,
2123
}
2224

23-
#[derive(Default)]
24-
pub(super) struct ReturnVisitor<'a> {
25+
pub(super) struct ReturnVisitor<'semantic, 'data> {
26+
/// The semantic model of the current file.
27+
semantic: &'semantic SemanticModel<'data>,
2528
/// The current stack of nodes.
26-
pub(super) stack: Stack<'a>,
29+
pub(super) stack: Stack<'data>,
2730
/// The preceding sibling of the current node.
28-
sibling: Option<&'a Stmt>,
31+
sibling: Option<&'data Stmt>,
2932
/// The parent nodes of the current node.
30-
parents: Vec<&'a Stmt>,
33+
parents: Vec<&'data Stmt>,
34+
}
35+
36+
impl<'semantic, 'data> ReturnVisitor<'semantic, 'data> {
37+
pub(super) fn new(semantic: &'semantic SemanticModel<'data>) -> Self {
38+
Self {
39+
semantic,
40+
stack: Stack::default(),
41+
sibling: None,
42+
parents: Vec::new(),
43+
}
44+
}
3145
}
3246

33-
impl<'a> Visitor<'a> for ReturnVisitor<'a> {
47+
impl<'semantic, 'a> Visitor<'a> for ReturnVisitor<'semantic, 'a> {
3448
fn visit_stmt(&mut self, stmt: &'a Stmt) {
3549
match stmt {
3650
Stmt::ClassDef(ast::StmtClassDef { decorator_list, .. }) => {
@@ -95,11 +109,17 @@ impl<'a> Visitor<'a> for ReturnVisitor<'a> {
95109
// x = f.read()
96110
// return x
97111
// ```
98-
Stmt::With(ast::StmtWith { body, .. }) => {
99-
if let Some(stmt_assign) = body.last().and_then(Stmt::as_assign_stmt) {
100-
self.stack
101-
.assignment_return
102-
.push((stmt_assign, stmt_return, stmt));
112+
Stmt::With(with) => {
113+
if let Some(stmt_assign) =
114+
with.body.last().and_then(Stmt::as_assign_stmt)
115+
{
116+
if !has_conditional_body(with, self.semantic) {
117+
self.stack.assignment_return.push((
118+
stmt_assign,
119+
stmt_return,
120+
stmt,
121+
));
122+
}
103123
}
104124
}
105125
_ => {}
@@ -142,3 +162,47 @@ impl<'a> Visitor<'a> for ReturnVisitor<'a> {
142162
self.sibling = sibling;
143163
}
144164
}
165+
166+
/// RET504
167+
/// If the last statement is a `return` statement, and the second-to-last statement is a
168+
/// `with` statement that suppresses an exception, then we should not analyze the `return`
169+
/// statement for unnecessary assignments. Otherwise we will suggest removing the assignment
170+
/// and the `with` statement, which would change the behavior of the code.
171+
///
172+
/// Example:
173+
/// ```python
174+
/// def foo(data):
175+
/// with suppress(JSONDecoderError):
176+
/// data = data.decode()
177+
/// return data
178+
179+
/// Returns `true` if the [`With`] statement is known to have a conditional body. In other words:
180+
/// if the [`With`] statement's body may or may not run.
181+
///
182+
/// For example, in the following, it's unsafe to inline the `return` into the `with`, since if
183+
/// `data.decode()` fails, the behavior of the program will differ. (As-is, the function will return
184+
/// the input `data`; if we inline the `return`, the function will return `None`.)
185+
///
186+
/// ```python
187+
/// def func(data):
188+
/// with suppress(JSONDecoderError):
189+
/// data = data.decode()
190+
/// return data
191+
/// ```
192+
fn has_conditional_body(with: &ast::StmtWith, semantic: &SemanticModel) -> bool {
193+
with.items.iter().any(|item| {
194+
let ast::WithItem {
195+
context_expr: Expr::Call(ast::ExprCall { func, .. }),
196+
..
197+
} = item
198+
else {
199+
return false;
200+
};
201+
if let Some(call_path) = semantic.resolve_call_path(func) {
202+
if call_path.as_slice() == ["contextlib", "suppress"] {
203+
return true;
204+
}
205+
}
206+
false
207+
})
208+
}

0 commit comments

Comments
 (0)