[flake8-return] Fix false-positive for variables used inside nested functions in RET504 (#18433)

<!--
Thank you for contributing to Ruff/ty! To help us out with reviewing,
please consider the following:

- Does this pull request include a summary of the change? (See below.)
- Does this pull request include a descriptive title? (Please prefix
with `[ty]` for ty pull
  requests.)
- Does this pull request include references to any relevant issues?
-->

## Summary

<!-- What's the purpose of the change? What does it do, and why? -->
This PR is the same as #17656.

I accidentally deleted the branch of that PR, so I'm creating a new one.

Fixes #14052

## Test Plan

Add regression tests
<!-- How was it tested? -->

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

@@ -422,6 +422,35 @@ def func(a: dict[str, int]) -> list[dict[str, int]]:
         services = a["services"]
         return services
 
+
+# See: https://github.com/astral-sh/ruff/issues/14052
+def outer() -> list[object]:
+    @register
+    async def inner() -> None:
+        print(layout)
+
+    layout = [...]
+    return layout
+
+def outer() -> list[object]:
+    with open("") as f:
+        async def inner() -> None:
+            print(layout)
+
+    layout = [...]
+    return layout
+
+
+def outer() -> list[object]:
+    def inner():
+        with open("") as f:
+            async def inner_inner() -> None:
+                print(layout)
+
+    layout = [...]
+    return layout
+
+
 # See: https://github.com/astral-sh/ruff/issues/18411
 def f():
     (#=

crates/ruff_linter/src/checkers/ast/analyze/bindings.rs

@@ -4,8 +4,8 @@ use crate::Fix;
 use crate::checkers::ast::Checker;
 use crate::codes::Rule;
 use crate::rules::{
-    flake8_import_conventions, flake8_pyi, flake8_pytest_style, flake8_type_checking, pyflakes,
-    pylint, pyupgrade, refurb, ruff,
+    flake8_import_conventions, flake8_pyi, flake8_pytest_style, flake8_return,
+    flake8_type_checking, pyflakes, pylint, pyupgrade, refurb, ruff,
 };
 
 /// Run lint rules over the [`Binding`]s.
@@ -25,11 +25,20 @@ pub(crate) fn bindings(checker: &Checker) {
         Rule::ForLoopWrites,
         Rule::CustomTypeVarForSelf,
         Rule::PrivateTypeParameter,
+        Rule::UnnecessaryAssign,
     ]) {
         return;
     }
 
     for (binding_id, binding) in checker.semantic.bindings.iter_enumerated() {
+        if checker.is_rule_enabled(Rule::UnnecessaryAssign) {
+            if binding.kind.is_function_definition() {
+                flake8_return::rules::unnecessary_assign(
+                    checker,
+                    binding.statement(checker.semantic()).unwrap(),
+                );
+            }
+        }
         if checker.is_rule_enabled(Rule::UnusedVariable) {
             if binding.kind.is_bound_exception()
                 && binding.is_unused()

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -207,7 +207,6 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                 Rule::UnnecessaryReturnNone,
                 Rule::ImplicitReturnValue,
                 Rule::ImplicitReturn,
-                Rule::UnnecessaryAssign,
                 Rule::SuperfluousElseReturn,
                 Rule::SuperfluousElseRaise,
                 Rule::SuperfluousElseContinue,

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

@@ -539,7 +539,21 @@ fn implicit_return(checker: &Checker, function_def: &ast::StmtFunctionDef, stmt:
 }
 
 /// RET504
-fn unnecessary_assign(checker: &Checker, stack: &Stack) {
+pub(crate) fn unnecessary_assign(checker: &Checker, function_stmt: &Stmt) {
+    let Stmt::FunctionDef(function_def) = function_stmt else {
+        return;
+    };
+    let Some(stack) = create_stack(checker, function_def) else {
+        return;
+    };
+
+    if !result_exists(&stack.returns) {
+        return;
+    }
+
+    let Some(function_scope) = checker.semantic().function_scope(function_def) else {
+        return;
+    };
     for (assign, return_, stmt) in &stack.assignment_return {
         // Identify, e.g., `return x`.
         let Some(value) = return_.value.as_ref() else {
@@ -583,6 +597,22 @@ fn unnecessary_assign(checker: &Checker, stack: &Stack) {
             continue;
         }
 
+        let Some(assigned_binding) = function_scope
+            .get(assigned_id)
+            .map(|binding_id| checker.semantic().binding(binding_id))
+        else {
+            continue;
+        };
+        // Check if there's any reference made to `assigned_binding` in another scope, e.g, nested
+        // functions. If there is, ignore them.
+        if assigned_binding
+            .references()
+            .map(|reference_id| checker.semantic().reference(reference_id))
+            .any(|reference| reference.scope_id() != assigned_binding.scope)
+        {
+            continue;
+        }
+
         let mut diagnostic = checker.report_diagnostic(
             UnnecessaryAssign {
                 name: assigned_id.to_string(),
@@ -665,24 +695,21 @@ fn superfluous_elif_else(checker: &Checker, stack: &Stack) {
     }
 }
 
-/// Run all checks from the `flake8-return` plugin.
-pub(crate) fn function(checker: &Checker, function_def: &ast::StmtFunctionDef) {
-    let ast::StmtFunctionDef {
-        decorator_list,
-        returns,
-        body,
-        ..
-    } = function_def;
+fn create_stack<'a>(
+    checker: &'a Checker,
+    function_def: &'a ast::StmtFunctionDef,
+) -> Option<Stack<'a>> {
+    let ast::StmtFunctionDef { body, .. } = function_def;
 
     // Find the last statement in the function.
     let Some(last_stmt) = body.last() else {
         // Skip empty functions.
-        return;
+        return None;
     };
 
     // Skip functions that consist of a single return statement.
     if body.len() == 1 && matches!(last_stmt, Stmt::Return(_)) {
-        return;
+        return None;
     }
 
     // Traverse the function body, to collect the stack.
@@ -696,9 +723,29 @@ pub(crate) fn function(checker: &Checker, function_def: &ast::StmtFunctionDef) {
 
     // Avoid false positives for generators.
     if stack.is_generator {
-        return;
+        return None;
     }
 
+    Some(stack)
+}
+
+/// Run all checks from the `flake8-return` plugin, but `RET504` which is ran
+/// after the semantic model is fully built.
+pub(crate) fn function(checker: &Checker, function_def: &ast::StmtFunctionDef) {
+    let ast::StmtFunctionDef {
+        decorator_list,
+        returns,
+        body,
+        ..
+    } = function_def;
+
+    let Some(stack) = create_stack(checker, function_def) else {
+        return;
+    };
+    let Some(last_stmt) = body.last() else {
+        return;
+    };
+
     if checker.any_rule_enabled(&[
         Rule::SuperfluousElseReturn,
         Rule::SuperfluousElseRaise,
@@ -721,10 +768,6 @@ pub(crate) fn function(checker: &Checker, function_def: &ast::StmtFunctionDef) {
         if checker.is_rule_enabled(Rule::ImplicitReturn) {
             implicit_return(checker, function_def, last_stmt);
         }
-
-        if checker.is_rule_enabled(Rule::UnnecessaryAssign) {
-            unnecessary_assign(checker, &stack);
-        }
     } else {
         if checker.is_rule_enabled(Rule::UnnecessaryReturnNone) {
             // Skip functions that have a return annotation that is not `None`.

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

@@ -247,8 +247,6 @@ RET504.py:423:16: RET504 [*] Unnecessary assignment to `services` before `return
 422 |         services = a["services"]
 423 |         return services
     |                ^^^^^^^^ RET504
-424 |
-425 | # See: https://github.com/astral-sh/ruff/issues/18411
     |
     = help: Remove unnecessary assignment
 
@@ -260,46 +258,46 @@ RET504.py:423:16: RET504 [*] Unnecessary assignment to `services` before `return
 423     |-        return services
     422 |+        return a["services"]
 424 423 | 
-425 424 | # See: https://github.com/astral-sh/ruff/issues/18411
-426 425 | def f():
+425 424 | 
+426 425 | # See: https://github.com/astral-sh/ruff/issues/14052
 
-RET504.py:429:12: RET504 [*] Unnecessary assignment to `x` before `return` statement
+RET504.py:458:12: RET504 [*] Unnecessary assignment to `x` before `return` statement
     |
-427 |     (#=
-428 |     x) = 1
-429 |     return x
+456 |     (#=
+457 |     x) = 1
+458 |     return x
     |            ^ RET504
-430 |
-431 | def f():
+459 |
+460 | def f():
     |
     = help: Remove unnecessary assignment
 
 ℹ Unsafe fix
-424 424 | 
-425 425 | # See: https://github.com/astral-sh/ruff/issues/18411
-426 426 | def f():
-427     |-    (#=
-428     |-    x) = 1
-429     |-    return x
-    427 |+    return 1
-430 428 | 
-431 429 | def f():
-432 430 |     x = (1
+453 453 | 
+454 454 | # See: https://github.com/astral-sh/ruff/issues/18411
+455 455 | def f():
+456     |-    (#=
+457     |-    x) = 1
+458     |-    return x
+    456 |+    return 1
+459 457 | 
+460 458 | def f():
+461 459 |     x = (1
 
-RET504.py:434:12: RET504 [*] Unnecessary assignment to `x` before `return` statement
+RET504.py:463:12: RET504 [*] Unnecessary assignment to `x` before `return` statement
     |
-432 |     x = (1
-433 |               )
-434 |     return x
+461 |     x = (1
+462 |               )
+463 |     return x
     |            ^ RET504
     |
     = help: Remove unnecessary assignment
 
 ℹ Unsafe fix
-429 429 |     return x
-430 430 | 
-431 431 | def f():
-432     |-    x = (1
-    432 |+    return (1
-433 433 |               )
-434     |-    return x
+458 458 |     return x
+459 459 | 
+460 460 | def f():
+461     |-    x = (1
+    461 |+    return (1
+462 462 |               )
+463     |-    return x

crates/ruff_python_semantic/src/model.rs

@@ -2094,6 +2094,20 @@ impl<'a> SemanticModel<'a> {
             None
         })
     }
+
+    /// Finds and returns the [`Scope`] corresponding to a given [`ast::StmtFunctionDef`].
+    ///
+    /// This method searches all scopes created by a function definition, comparing the
+    /// [`TextRange`] of the provided `function_def` with the the range of the function
+    /// associated with the scope.
+    pub fn function_scope(&self, function_def: &ast::StmtFunctionDef) -> Option<&Scope> {
+        self.scopes.iter().find(|scope| {
+            let Some(function) = scope.kind.as_function() else {
+                return false;
+            };
+            function.range() == function_def.range()
+        })
+    }
 }
 
 pub struct ShadowedBinding {