Make some `flake8-return` rules auto-fixable (#1017)

README.md

@@ -728,9 +728,9 @@ For more, see [flake8-return](https://pypi.org/project/flake8-return/1.2.0/) on
 
 | Code | Name | Message | Fix |
 | ---- | ---- | ------- | --- |
-| RET501 | UnnecessaryReturnNone | Do not explicitly `return None` in function if it is the only possible return value |  |
+| RET501 | UnnecessaryReturnNone | Do not explicitly `return None` in function if it is the only possible return value | 🛠 |
-| RET502 | ImplicitReturnValue | Do not implicitly `return None` in function able to return non-`None` value |  |
+| RET502 | ImplicitReturnValue | Do not implicitly `return None` in function able to return non-`None` value | 🛠 |
-| RET503 | ImplicitReturn | Missing explicit `return` at the end of function able to return non-`None` value |  |
+| RET503 | ImplicitReturn | Missing explicit `return` at the end of function able to return non-`None` value | 🛠 |
 | RET504 | UnnecessaryAssign | Unnecessary variable assignment before `return` statement |  |
 | RET505 | SuperfluousElseReturn | Unnecessary `else` after `return` statement |  |
 | RET506 | SuperfluousElseRaise | Unnecessary `else` after `raise` statement |  |

src/checks.rs

@@ -2469,6 +2469,8 @@ impl CheckKind {
                 | CheckKind::DoNotAssignLambda
                 | CheckKind::DuplicateHandlerException(..)
                 | CheckKind::GetAttrWithConstant
+                | CheckKind::ImplicitReturn
+                | CheckKind::ImplicitReturnValue
                 | CheckKind::IsLiteral
                 | CheckKind::NewLineAfterLastParagraph
                 | CheckKind::NewLineAfterSectionName(..)
@@ -2488,6 +2490,7 @@ impl CheckKind {
                 | CheckKind::PPrintFound
                 | CheckKind::PrintFound
                 | CheckKind::RaiseNotImplemented
+                | CheckKind::RedundantOpenModes
                 | CheckKind::RedundantTupleInExceptionHandler(..)
                 | CheckKind::SectionNameEndsInColon(..)
                 | CheckKind::SectionNotOverIndented(..)
@@ -2502,7 +2505,6 @@ impl CheckKind {
                 | CheckKind::UnnecessaryComprehension(..)
                 | CheckKind::UnnecessaryEncodeUTF8
                 | CheckKind::UnnecessaryFutureImport(..)
-                | CheckKind::RedundantOpenModes
                 | CheckKind::UnnecessaryGeneratorDict
                 | CheckKind::UnnecessaryGeneratorList
                 | CheckKind::UnnecessaryGeneratorSet
@@ -2514,6 +2516,7 @@ impl CheckKind {
                 | CheckKind::UnnecessaryLiteralSet(..)
                 | CheckKind::UnnecessaryLiteralWithinListCall(..)
                 | CheckKind::UnnecessaryLiteralWithinTupleCall(..)
+                | CheckKind::UnnecessaryReturnNone
                 | CheckKind::UnsortedImports
                 | CheckKind::UnusedImport(_, false)
                 | CheckKind::UnusedLoopControlVariable(..)

src/flake8_return/plugins.rs

@@ -3,6 +3,8 @@ use rustpython_ast::{Constant, Expr, ExprKind, Location, Stmt, StmtKind};
 
 use crate::ast::types::Range;
 use crate::ast::visitor::Visitor;
+use crate::ast::whitespace::indentation;
+use crate::autofix::Fix;
 use crate::check_ast::Checker;
 use crate::checks::{Branch, CheckCode, CheckKind};
 use crate::flake8_return::helpers::result_exists;
@@ -20,10 +22,16 @@ fn unnecessary_return_none(checker: &mut Checker, stack: &Stack) {
                     ..
                 }
             ) {
-                checker.add_check(Check::new(
+                let mut check =
-                    CheckKind::UnnecessaryReturnNone,
+                    Check::new(CheckKind::UnnecessaryReturnNone, Range::from_located(stmt));
-                    Range::from_located(stmt),
+                if checker.patch(&CheckCode::RET501) {
-                ));
+                    check.amend(Fix::replacement(
+                        "return".to_string(),
+                        stmt.location,
+                        stmt.end_location.unwrap(),
+                    ));
+                }
+                checker.add_check(check);
             }
         }
     }
@@ -33,10 +41,15 @@ fn unnecessary_return_none(checker: &mut Checker, stack: &Stack) {
 fn implicit_return_value(checker: &mut Checker, stack: &Stack) {
     for (stmt, expr) in &stack.returns {
         if expr.is_none() {
-            checker.add_check(Check::new(
+            let mut check = Check::new(CheckKind::ImplicitReturnValue, Range::from_located(stmt));
-                CheckKind::ImplicitReturnValue,
+            if checker.patch(&CheckCode::RET502) {
-                Range::from_located(stmt),
+                check.amend(Fix::replacement(
-            ));
+                    "return None".to_string(),
+                    stmt.location,
+                    stmt.end_location.unwrap(),
+                ));
+            }
+            checker.add_check(check);
         }
     }
 }
@@ -85,10 +98,18 @@ fn implicit_return(checker: &mut Checker, last_stmt: &Stmt) {
         | StmtKind::Raise { .. }
         | StmtKind::Try { .. } => {}
         _ => {
-            checker.add_check(Check::new(
+            let mut check = Check::new(CheckKind::ImplicitReturn, Range::from_located(last_stmt));
-                CheckKind::ImplicitReturn,
+            if checker.patch(&CheckCode::RET503) {
-                Range::from_located(last_stmt),
+                let mut content = String::new();
-            ));
+                content.push_str(&indentation(checker, last_stmt));
+                content.push_str("return None");
+                content.push('\n');
+                check.amend(Fix::insertion(
+                    content,
+                    Location::new(last_stmt.end_location.unwrap().row() + 1, 0),
+                ));
+            }
+            checker.add_check(check);
         }
     }
 }

src/flake8_return/snapshots/ruff__flake8_return__tests__RET501_RET501.py.snap

@@ -9,5 +9,12 @@ expression: checks
   end_location:
     row: 4
     column: 15
-  fix: ~
+  fix:
+    content: return
+    location:
+      row: 4
+      column: 4
+    end_location:
+      row: 4
+      column: 15
 

src/flake8_return/snapshots/ruff__flake8_return__tests__RET502_RET502.py.snap

@@ -9,5 +9,12 @@ expression: checks
   end_location:
     row: 3
     column: 14
-  fix: ~
+  fix:
+    content: return None
+    location:
+      row: 3
+      column: 8
+    end_location:
+      row: 3
+      column: 14
 

src/flake8_return/snapshots/ruff__flake8_return__tests__RET503_RET503.py.snap

@@ -17,7 +17,14 @@ expression: checks
   end_location:
     row: 14
     column: 15
-  fix: ~
+  fix:
+    content: "        return None\n"
+    location:
+      row: 15
+      column: 0
+    end_location:
+      row: 15
+      column: 0
 - kind: ImplicitReturn
   location:
     row: 23
@@ -25,7 +32,14 @@ expression: checks
   end_location:
     row: 23
     column: 11
-  fix: ~
+  fix:
+    content: "    return None\n"
+    location:
+      row: 24
+      column: 0
+    end_location:
+      row: 24
+      column: 0
 - kind: ImplicitReturn
   location:
     row: 29
@@ -41,5 +55,12 @@ expression: checks
   end_location:
     row: 39
     column: 15
-  fix: ~
+  fix:
+    content: "        return None\n"
+    location:
+      row: 40
+      column: 0
+    end_location:
+      row: 40
+      column: 0