[flake8-return] - Add fixes for (RET505, RET506, RET507, `RET…

…508`) (#9595)

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

@@ -131,6 +131,37 @@ def bar3(x, y, z):
     return None
 
 
+def bar4(x):
+    if True:
+        return
+    else:
+        # comment
+        pass
+
+
+def bar5():
+    if True:
+        return
+    else:  # comment
+        pass
+
+
+def bar6():
+    if True:
+        return
+    else\
+        :\
+        # comment
+        pass
+
+
+def bar7():
+    if True:
+        return
+    else\
+        :  # comment
+        pass
+
 x = 0
 
 if x == 1:

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

@@ -11,10 +11,11 @@ mod tests {
     use anyhow::Result;
     use test_case::test_case;
 
-    use crate::assert_messages;
     use crate::registry::Rule;
+    use crate::settings::types::PreviewMode;
     use crate::settings::LinterSettings;
     use crate::test::test_path;
+    use crate::{assert_messages, settings};
 
     #[test_case(Rule::UnnecessaryReturnNone, Path::new("RET501.py"))]
     #[test_case(Rule::ImplicitReturnValue, Path::new("RET502.py"))]
@@ -33,4 +34,25 @@ mod tests {
         assert_messages!(snapshot, diagnostics);
         Ok(())
     }
+
+    #[test_case(Rule::SuperfluousElseReturn, Path::new("RET505.py"))]
+    #[test_case(Rule::SuperfluousElseRaise, Path::new("RET506.py"))]
+    #[test_case(Rule::SuperfluousElseContinue, Path::new("RET507.py"))]
+    #[test_case(Rule::SuperfluousElseBreak, Path::new("RET508.py"))]
+    fn preview_rules(rule_code: Rule, path: &Path) -> Result<()> {
+        let snapshot = format!(
+            "preview__{}_{}",
+            rule_code.noqa_code(),
+            path.to_string_lossy()
+        );
+        let diagnostics = test_path(
+            Path::new("flake8_return").join(path).as_path(),
+            &settings::LinterSettings {
+                preview: PreviewMode::Enabled,
+                ..settings::LinterSettings::for_rule(rule_code)
+            },
+        )?;
+        assert_messages!(snapshot, diagnostics);
+        Ok(())
+    }
 }

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

@@ -1,9 +1,10 @@
+use anyhow::Result;
 use std::ops::Add;
 
 use ruff_python_ast::{self as ast, ElifElseClause, Expr, Stmt};
 use ruff_text_size::{Ranged, TextRange, TextSize};
 
-use ruff_diagnostics::{AlwaysFixableViolation, Violation};
+use ruff_diagnostics::{AlwaysFixableViolation, FixAvailability, Violation};
 use ruff_diagnostics::{Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
 
@@ -12,12 +13,13 @@ use ruff_python_ast::stmt_if::elif_else_range;
 use ruff_python_ast::visitor::Visitor;
 use ruff_python_ast::whitespace::indentation;
 use ruff_python_semantic::SemanticModel;
-use ruff_python_trivia::is_python_whitespace;
+use ruff_python_trivia::{is_python_whitespace, SimpleTokenKind, SimpleTokenizer};
 
 use crate::checkers::ast::Checker;
 use crate::fix::edits;
 use crate::registry::{AsRule, Rule};
 use crate::rules::flake8_return::helpers::end_of_last_statement;
+use crate::rules::pyupgrade::fixes::adjust_indentation;
 
 use super::super::branch::Branch;
 use super::super::helpers::result_exists;
@@ -210,11 +212,17 @@ pub struct SuperfluousElseReturn {
 }
 
 impl Violation for SuperfluousElseReturn {
+    const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
     #[derive_message_formats]
     fn message(&self) -> String {
         let SuperfluousElseReturn { branch } = self;
         format!("Unnecessary `{branch}` after `return` statement")
     }
+
+    fn fix_title(&self) -> Option<String> {
+        let SuperfluousElseReturn { branch } = self;
+        Some(format!("Remove unnecessary `{branch}`"))
+    }
 }
 
 /// ## What it does
@@ -248,11 +256,17 @@ pub struct SuperfluousElseRaise {
 }
 
 impl Violation for SuperfluousElseRaise {
+    const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
     #[derive_message_formats]
     fn message(&self) -> String {
         let SuperfluousElseRaise { branch } = self;
         format!("Unnecessary `{branch}` after `raise` statement")
     }
+
+    fn fix_title(&self) -> Option<String> {
+        let SuperfluousElseRaise { branch } = self;
+        Some(format!("Remove unnecessary `{branch}`"))
+    }
 }
 
 /// ## What it does
@@ -288,11 +302,17 @@ pub struct SuperfluousElseContinue {
 }
 
 impl Violation for SuperfluousElseContinue {
+    const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
     #[derive_message_formats]
     fn message(&self) -> String {
         let SuperfluousElseContinue { branch } = self;
         format!("Unnecessary `{branch}` after `continue` statement")
     }
+
+    fn fix_title(&self) -> Option<String> {
+        let SuperfluousElseContinue { branch } = self;
+        Some(format!("Remove unnecessary `{branch}`"))
+    }
 }
 
 /// ## What it does
@@ -328,11 +348,17 @@ pub struct SuperfluousElseBreak {
 }
 
 impl Violation for SuperfluousElseBreak {
+    const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
     #[derive_message_formats]
     fn message(&self) -> String {
         let SuperfluousElseBreak { branch } = self;
         format!("Unnecessary `{branch}` after `break` statement")
     }
+
+    fn fix_title(&self) -> Option<String> {
+        let SuperfluousElseBreak { branch } = self;
+        Some(format!("Remove unnecessary `{branch}`"))
+    }
 }
 
 /// RET501
@@ -602,42 +628,54 @@ fn superfluous_else_node(
     };
     for child in if_elif_body {
         if child.is_return_stmt() {
-            let diagnostic = Diagnostic::new(
+            let mut diagnostic = Diagnostic::new(
                 SuperfluousElseReturn { branch },
                 elif_else_range(elif_else, checker.locator().contents())
                     .unwrap_or_else(|| elif_else.range()),
             );
             if checker.enabled(diagnostic.kind.rule()) {
+                if checker.settings.preview.is_enabled() {
+                    diagnostic.try_set_fix(|| remove_else(checker, elif_else));
+                }
                 checker.diagnostics.push(diagnostic);
             }
             return true;
         } else if child.is_break_stmt() {
-            let diagnostic = Diagnostic::new(
+            let mut diagnostic = Diagnostic::new(
                 SuperfluousElseBreak { branch },
                 elif_else_range(elif_else, checker.locator().contents())
                     .unwrap_or_else(|| elif_else.range()),
             );
             if checker.enabled(diagnostic.kind.rule()) {
+                if checker.settings.preview.is_enabled() {
+                    diagnostic.try_set_fix(|| remove_else(checker, elif_else));
+                }
                 checker.diagnostics.push(diagnostic);
             }
             return true;
         } else if child.is_raise_stmt() {
-            let diagnostic = Diagnostic::new(
+            let mut diagnostic = Diagnostic::new(
                 SuperfluousElseRaise { branch },
                 elif_else_range(elif_else, checker.locator().contents())
                     .unwrap_or_else(|| elif_else.range()),
             );
             if checker.enabled(diagnostic.kind.rule()) {
+                if checker.settings.preview.is_enabled() {
+                    diagnostic.try_set_fix(|| remove_else(checker, elif_else));
+                }
                 checker.diagnostics.push(diagnostic);
             }
             return true;
         } else if child.is_continue_stmt() {
-            let diagnostic = Diagnostic::new(
+            let mut diagnostic = Diagnostic::new(
                 SuperfluousElseContinue { branch },
                 elif_else_range(elif_else, checker.locator().contents())
                     .unwrap_or_else(|| elif_else.range()),
             );
             if checker.enabled(diagnostic.kind.rule()) {
+                if checker.settings.preview.is_enabled() {
+                    diagnostic.try_set_fix(|| remove_else(checker, elif_else));
+                }
                 checker.diagnostics.push(diagnostic);
             }
             return true;
@@ -715,3 +753,67 @@ pub(crate) fn function(checker: &mut Checker, body: &[Stmt], returns: Option<&Ex
         }
     }
 }
+
+fn remove_else(checker: &mut Checker, elif_else: &ElifElseClause) -> Result<Fix> {
+    if elif_else.test.is_some() {
+        // it's an elif, so we can just make it an if
+
+        // delete "el" from "elif"
+        Ok(Fix::safe_edit(Edit::deletion(
+            elif_else.start(),
+            elif_else.start() + TextSize::from(2),
+        )))
+    } else {
+        // the start of the line where the `else`` is
+        let else_line_start = checker.locator().line_start(elif_else.start());
+
+        // making a tokenizer to find the Colon for the `else`, not always on the same line!
+        let mut else_line_tokenizer =
+            SimpleTokenizer::starts_at(elif_else.start(), checker.locator().contents());
+
+        // find the Colon for the `else`
+        let Some(else_colon) =
+            else_line_tokenizer.find(|token| token.kind == SimpleTokenKind::Colon)
+        else {
+            return Err(anyhow::anyhow!("Cannot find `:` in `else` statement"));
+        };
+
+        // get the indentation of the `else`, since that is the indent level we want to end with
+        let Some(desired_indentation) = indentation(checker.locator(), elif_else) else {
+            return Err(anyhow::anyhow!("Compound statement cannot be inlined"));
+        };
+
+        // we're deleting the `else`, and it's Colon, and the rest of the line(s) they're on,
+        // so here we get the last position of the line the Colon is on
+        let else_colon_end = checker.locator().full_line_end(else_colon.end());
+
+        // if there is a comment on the same line as the Colon, let's keep it
+        // and give it the proper indentation once we unindent it
+        let else_comment_after_colon = else_line_tokenizer
+            .find(|token| token.kind.is_comment())
+            .and_then(|token| {
+                if token.kind == SimpleTokenKind::Comment && token.start() < else_colon_end {
+                    return Some(format!(
+                        "{desired_indentation}{}{}",
+                        checker.locator().slice(token),
+                        checker.stylist().line_ending().as_str(),
+                    ));
+                }
+                None
+            })
+            .unwrap_or(String::new());
+
+        let indented = adjust_indentation(
+            TextRange::new(else_colon_end, elif_else.end()),
+            desired_indentation,
+            checker.locator(),
+            checker.stylist(),
+        )?;
+
+        Ok(Fix::safe_edit(Edit::replacement(
+            format!("{else_comment_after_colon}{indented}"),
+            else_line_start,
+            elif_else.end(),
+        )))
+    }
+}