[flake8-return] Only add return None at end of function (RET503) (#11074

)

Co-authored-by: Micha Reiser <micha@reiser.io>

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

@@ -368,3 +368,11 @@ def bar() -> NoReturn:
     if baz() > 3:
         return 1
     bar()
+
+
+def f():
+    if a:
+        return b
+    else:
+        with c:
+            d

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

@@ -229,12 +229,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                 Rule::SuperfluousElseContinue,
                 Rule::SuperfluousElseBreak,
             ]) {
-                flake8_return::rules::function(
-                    checker,
-                    body,
-                    decorator_list,
-                    returns.as_ref().map(AsRef::as_ref),
-                );
+                flake8_return::rules::function(checker, function_def);
             }
             if checker.enabled(Rule::UselessReturn) {
                 pylint::rules::useless_return(

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

@@ -35,6 +35,7 @@ mod tests {
         Ok(())
     }
 
+    #[test_case(Rule::ImplicitReturn, Path::new("RET503.py"))]
     #[test_case(Rule::SuperfluousElseReturn, Path::new("RET505.py"))]
     #[test_case(Rule::SuperfluousElseRaise, Path::new("RET506.py"))]
     #[test_case(Rule::SuperfluousElseContinue, Path::new("RET507.py"))]

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

@@ -451,94 +451,110 @@ fn is_noreturn_func(func: &Expr, semantic: &SemanticModel) -> bool {
     semantic.match_typing_qualified_name(&qualified_name, "NoReturn")
 }
 
-/// RET503
-fn implicit_return(checker: &mut Checker, stmt: &Stmt) {
+fn add_return_none(checker: &mut Checker, stmt: &Stmt, range: TextRange) {
+    let mut diagnostic = Diagnostic::new(ImplicitReturn, range);
+    if let Some(indent) = indentation(checker.locator(), stmt) {
+        let mut content = String::new();
+        content.push_str(checker.stylist().line_ending().as_str());
+        content.push_str(indent);
+        content.push_str("return None");
+        diagnostic.set_fix(Fix::unsafe_edit(Edit::insertion(
+            content,
+            end_of_last_statement(stmt, checker.locator()),
+        )));
+    }
+    checker.diagnostics.push(diagnostic);
+}
+
+/// Returns a list of all implicit returns in the given statement.
+///
+/// Note: The function should be refactored to `has_implicit_return` with an early return (when seeing the first implicit return)
+/// when removing the preview gating.
+fn implicit_returns<'a>(checker: &Checker, stmt: &'a Stmt) -> Vec<&'a Stmt> {
     match stmt {
         Stmt::If(ast::StmtIf {
             body,
             elif_else_clauses,
             ..
         }) => {
-            if let Some(last_stmt) = body.last() {
-                implicit_return(checker, last_stmt);
-            }
+            let mut implicit_stmts = body
+                .last()
+                .map(|last| implicit_returns(checker, last))
+                .unwrap_or_default();
+
             for clause in elif_else_clauses {
-                if let Some(last_stmt) = clause.body.last() {
-                    implicit_return(checker, last_stmt);
-                }
+                implicit_stmts.extend(
+                    clause
+                        .body
+                        .last()
+                        .iter()
+                        .flat_map(|last| implicit_returns(checker, last)),
+                );
             }
 
             // Check if we don't have an else clause
             if matches!(
                 elif_else_clauses.last(),
                 None | Some(ast::ElifElseClause { test: Some(_), .. })
             ) {
-                let mut diagnostic = Diagnostic::new(ImplicitReturn, stmt.range());
-                if let Some(indent) = indentation(checker.locator(), stmt) {
-                    let mut content = String::new();
-                    content.push_str(checker.stylist().line_ending().as_str());
-                    content.push_str(indent);
-                    content.push_str("return None");
-                    diagnostic.set_fix(Fix::unsafe_edit(Edit::insertion(
-                        content,
-                        end_of_last_statement(stmt, checker.locator()),
-                    )));
-                }
-                checker.diagnostics.push(diagnostic);
+                implicit_stmts.push(stmt);
             }
+            implicit_stmts
         }
-        Stmt::Assert(ast::StmtAssert { test, .. }) if is_const_false(test) => {}
-        Stmt::While(ast::StmtWhile { test, .. }) if is_const_true(test) => {}
+        Stmt::Assert(ast::StmtAssert { test, .. }) if is_const_false(test) => vec![],
+        Stmt::While(ast::StmtWhile { test, .. }) if is_const_true(test) => vec![],
         Stmt::For(ast::StmtFor { orelse, .. }) | Stmt::While(ast::StmtWhile { orelse, .. }) => {
             if let Some(last_stmt) = orelse.last() {
-                implicit_return(checker, last_stmt);
+                implicit_returns(checker, last_stmt)
             } else {
-                let mut diagnostic = Diagnostic::new(ImplicitReturn, stmt.range());
-                if let Some(indent) = indentation(checker.locator(), stmt) {
-                    let mut content = String::new();
-                    content.push_str(checker.stylist().line_ending().as_str());
-                    content.push_str(indent);
-                    content.push_str("return None");
-                    diagnostic.set_fix(Fix::unsafe_edit(Edit::insertion(
-                        content,
-                        end_of_last_statement(stmt, checker.locator()),
-                    )));
-                }
-                checker.diagnostics.push(diagnostic);
+                vec![stmt]
             }
         }
         Stmt::Match(ast::StmtMatch { cases, .. }) => {
+            let mut implicit_stmts = vec![];
             for case in cases {
-                if let Some(last_stmt) = case.body.last() {
-                    implicit_return(checker, last_stmt);
-                }
-            }
-        }
-        Stmt::With(ast::StmtWith { body, .. }) => {
-            if let Some(last_stmt) = body.last() {
-                implicit_return(checker, last_stmt);
+                implicit_stmts.extend(
+                    case.body
+                        .last()
+                        .into_iter()
+                        .flat_map(|last_stmt| implicit_returns(checker, last_stmt)),
+                );
             }
+            implicit_stmts
         }
-        Stmt::Return(_) | Stmt::Raise(_) | Stmt::Try(_) => {}
+        Stmt::With(ast::StmtWith { body, .. }) => body
+            .last()
+            .map(|last_stmt| implicit_returns(checker, last_stmt))
+            .unwrap_or_default(),
+        Stmt::Return(_) | Stmt::Raise(_) | Stmt::Try(_) => vec![],
         Stmt::Expr(ast::StmtExpr { value, .. })
             if matches!(
                 value.as_ref(),
                 Expr::Call(ast::ExprCall { func, ..  })
                     if is_noreturn_func(func, checker.semantic())
-            ) => {}
+            ) =>
+        {
+            vec![]
+        }
         _ => {
-            let mut diagnostic = Diagnostic::new(ImplicitReturn, stmt.range());
-            if let Some(indent) = indentation(checker.locator(), stmt) {
-                let mut content = String::new();
-                content.push_str(checker.stylist().line_ending().as_str());
-                content.push_str(indent);
-                content.push_str("return None");
-                diagnostic.set_fix(Fix::unsafe_edit(Edit::insertion(
-                    content,
-                    end_of_last_statement(stmt, checker.locator()),
-                )));
-            }
-            checker.diagnostics.push(diagnostic);
+            vec![stmt]
+        }
+    }
+}
+
+/// RET503
+fn implicit_return(checker: &mut Checker, function_def: &ast::StmtFunctionDef, stmt: &Stmt) {
+    let implicit_stmts = implicit_returns(checker, stmt);
+
+    if implicit_stmts.is_empty() {
+        return;
+    }
+
+    if checker.settings.preview.is_enabled() {
+        add_return_none(checker, stmt, function_def.range());
+    } else {
+        for implicit_stmt in implicit_stmts {
+            add_return_none(checker, implicit_stmt, implicit_stmt.range());
         }
     }
 }
@@ -742,12 +758,14 @@ fn superfluous_elif_else(checker: &mut Checker, stack: &Stack) {
 }
 
 /// Run all checks from the `flake8-return` plugin.
-pub(crate) fn function(
-    checker: &mut Checker,
-    body: &[Stmt],
-    decorator_list: &[Decorator],
-    returns: Option<&Expr>,
-) {
+pub(crate) fn function(checker: &mut Checker, function_def: &ast::StmtFunctionDef) {
+    let ast::StmtFunctionDef {
+        decorator_list,
+        returns,
+        body,
+        ..
+    } = function_def;
+
     // Find the last statement in the function.
     let Some(last_stmt) = body.last() else {
         // Skip empty functions.
@@ -793,7 +811,7 @@ pub(crate) fn function(
             implicit_return_value(checker, &stack);
         }
         if checker.enabled(Rule::ImplicitReturn) {
-            implicit_return(checker, last_stmt);
+            implicit_return(checker, function_def, last_stmt);
         }
 
         if checker.enabled(Rule::UnnecessaryAssign) {
@@ -802,7 +820,7 @@ pub(crate) fn function(
     } else {
         if checker.enabled(Rule::UnnecessaryReturnNone) {
             // Skip functions that have a return annotation that is not `None`.
-            if returns.map_or(true, Expr::is_none_literal_expr) {
+            if returns.as_deref().map_or(true, Expr::is_none_literal_expr) {
                 unnecessary_return_none(checker, decorator_list, &stack);
             }
         }

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

@@ -452,5 +452,21 @@ RET503.py:370:5: RET503 [*] Missing explicit `return` at the end of function abl
 369 369 |         return 1
 370 370 |     bar()
     371 |+    return None
+371 372 | 
+372 373 | 
+373 374 | def f():
 
+RET503.py:378:13: RET503 [*] Missing explicit `return` at the end of function able to return non-`None` value
+    |
+376 |     else:
+377 |         with c:
+378 |             d
+    |             ^ RET503
+    |
+    = help: Add explicit `return` statement
 
+ℹ Unsafe fix
+376 376 |     else:
+377 377 |         with c:
+378 378 |             d
+    379 |+            return None