From 9fa0d35cb5ecd4457e2403e5ba7b7fda98a10569 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Tue, 2 Apr 2024 15:50:04 -0400 Subject: [PATCH] Ignore assignments to annotated variables in unnecessary-assign --- .../test/fixtures/flake8_return/RET504.py | 15 ++++++++ .../src/rules/flake8_return/rules/function.rs | 6 ++++ ...lake8_return__tests__RET504_RET504.py.snap | 15 ++++++++ .../src/rules/flake8_return/visitor.rs | 36 ++++++++++++------- 4 files changed, 59 insertions(+), 13 deletions(-) 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 8480aafaa1c01..91a60a7540dc9 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_return/RET504.py +++ b/crates/ruff_linter/resources/test/fixtures/flake8_return/RET504.py @@ -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 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 9e0f90c4077e4..124abdfee953e 100644 --- a/crates/ruff_linter/src/rules/flake8_return/rules/function.rs +++ b/crates/ruff_linter/src/rules/flake8_return/rules/function.rs @@ -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; } 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 3656e71afbd2d..46ffb2df68b29 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 @@ -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"] diff --git a/crates/ruff_linter/src/rules/flake8_return/visitor.rs b/crates/ruff_linter/src/rules/flake8_return/visitor.rs index f583c22d6cab0..c88abc6676bd5 100644 --- a/crates/ruff_linter/src/rules/flake8_return/visitor.rs +++ b/crates/ruff_linter/src/rules/flake8_return/visitor.rs @@ -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. @@ -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. @@ -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. ///