Skip to content

Commit

Permalink
[flake8-return] Ignore assignments to annotated variables in `unnec…
Browse files Browse the repository at this point in the history
…essary-assign` (#10741)

## Summary

Closes #10732.
  • Loading branch information
charliermarsh committed Apr 2, 2024
1 parent 859e3fc commit dff8f93
Show file tree
Hide file tree
Showing 4 changed files with 59 additions and 13 deletions.
15 changes: 15 additions & 0 deletions crates/ruff_linter/resources/test/fixtures/flake8_return/RET504.py
Original file line number Diff line number Diff line change
Expand Up @@ -406,3 +406,18 @@ def foo():
with contextlib.suppress(Exception):
y = 2
return y


# See: https://github.com/astral-sh/ruff/issues/10732
def func(a: dict[str, int]) -> list[dict[str, int]]:
services: list[dict[str, int]]
if "services" in a:
services = a["services"]
return services


# See: https://github.com/astral-sh/ruff/issues/10732
def func(a: dict[str, int]) -> list[dict[str, int]]:
if "services" in a:
services = a["services"]
return services
6 changes: 6 additions & 0 deletions crates/ruff_linter/src/rules/flake8_return/rules/function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -567,6 +567,12 @@ fn unnecessary_assign(checker: &mut Checker, stack: &Stack) {
continue;
}

// Ignore variables that have an annotation defined elsewhere.
if stack.annotations.contains(assigned_id.as_str()) {
continue;
}

// Ignore `nonlocal` and `global` variables.
if stack.non_locals.contains(assigned_id.as_str()) {
continue;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -241,4 +241,19 @@ RET504.py:400:12: RET504 [*] Unnecessary assignment to `y` before `return` state
402 401 |
403 402 | def foo():

RET504.py:423:16: RET504 [*] Unnecessary assignment to `services` before `return` statement
|
421 | if "services" in a:
422 | services = a["services"]
423 | return services
| ^^^^^^^^ RET504
|
= help: Remove unnecessary assignment

Unsafe fix
419 419 | # See: https://github.com/astral-sh/ruff/issues/10732
420 420 | def func(a: dict[str, int]) -> list[dict[str, int]]:
421 421 | if "services" in a:
422 |- services = a["services"]
423 |- return services
422 |+ return a["services"]
36 changes: 23 additions & 13 deletions crates/ruff_linter/src/rules/flake8_return/visitor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,21 @@ pub(super) struct Stack<'data> {
pub(super) elifs_elses: Vec<(&'data [Stmt], &'data ElifElseClause)>,
/// The non-local variables in the current function.
pub(super) non_locals: FxHashSet<&'data str>,
/// The annotated variables in the current function.
///
/// For example, consider:
/// ```python
/// x: int
///
/// if True:
/// x = foo()
/// return x
/// ```
///
/// In this case, the annotation on `x` is used to cast the return value
/// of `foo()` to an `int`. Removing the `x = foo()` statement would
/// change the return type of the function.
pub(super) annotations: 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.
Expand Down Expand Up @@ -86,6 +101,14 @@ impl<'semantic, 'a> Visitor<'a> for ReturnVisitor<'semantic, 'a> {
.non_locals
.extend(names.iter().map(Identifier::as_str));
}
Stmt::AnnAssign(ast::StmtAnnAssign { target, value, .. }) => {
// Ex) `x: int`
if value.is_none() {
if let Expr::Name(name) = target.as_ref() {
self.stack.annotations.insert(name.id.as_str());
}
}
}
Stmt::Return(stmt_return) => {
// If the `return` statement is preceded by an `assignment` statement, then the
// `assignment` statement may be redundant.
Expand Down Expand Up @@ -163,19 +186,6 @@ impl<'semantic, 'a> Visitor<'a> for ReturnVisitor<'semantic, 'a> {
}
}

/// 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.
///
Expand Down

0 comments on commit dff8f93

Please sign in to comment.