keep comments after else colon + add a couple more tests

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/rules/function.rs

@@ -13,7 +13,7 @@ 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;
@@ -737,22 +737,55 @@ fn remove_else(checker: &mut Checker, elif_else: &ElifElseClause) -> Result<Fix>
             elif_else.start() + TextSize::from(2),
         )))
     } else {
-        let else_line_range = checker.locator().full_line_range(elif_else.start());
+        // 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_line_range.end(), elif_else.end()),
+            TextRange::new(else_colon_end, elif_else.end()),
             desired_indentation,
             checker.locator(),
             checker.stylist(),
         )?;
 
         Ok(Fix::safe_edit(Edit::replacement(
-            indented,
-            else_line_range.start(),
+            format!("{else_comment_after_colon}{indented}"),
+            else_line_start,
             elif_else.end(),
         )))
     }

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

@@ -89,4 +89,47 @@ RET505.py:99:5: RET505 Unnecessary `else` after `return` statement
     |
     = help: Remove unnecessary `else`
 
+RET505.py:137:5: RET505 Unnecessary `else` after `return` statement
+    |
+135 |     if True:
+136 |         return
+137 |     else:
+    |     ^^^^ RET505
+138 |         # comment
+139 |         pass
+    |
+    = help: Remove unnecessary `else`
+
+RET505.py:145:5: RET505 Unnecessary `else` after `return` statement
+    |
+143 |     if True:
+144 |         return
+145 |     else:  # comment
+    |     ^^^^ RET505
+146 |         pass
+    |
+    = help: Remove unnecessary `else`
+
+RET505.py:152:5: RET505 Unnecessary `else` after `return` statement
+    |
+150 |     if True:
+151 |         return
+152 |     else\
+    |     ^^^^ RET505
+153 |         :\
+154 |         # comment
+    |
+    = help: Remove unnecessary `else`
+
+RET505.py:161:5: RET505 Unnecessary `else` after `return` statement
+    |
+159 |     if True:
+160 |         return
+161 |     else\
+    |     ^^^^ RET505
+162 |         :  # comment
+163 |         pass
+    |
+    = help: Remove unnecessary `else`
+
 

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

@@ -184,4 +184,99 @@ RET505.py:99:5: RET505 [*] Unnecessary `else` after `return` statement
 105 104 | 
 106 105 | ###
 
+RET505.py:137:5: RET505 [*] Unnecessary `else` after `return` statement
+    |
+135 |     if True:
+136 |         return
+137 |     else:
+    |     ^^^^ RET505
+138 |         # comment
+139 |         pass
+    |
+    = help: Remove unnecessary `else`
+
+ℹ Safe fix
+134 134 | def bar4(x):
+135 135 |     if True:
+136 136 |         return
+137     |-    else:
+138     |-        # comment
+139     |-        pass
+    137 |+    # comment
+    138 |+    pass
+140 139 | 
+141 140 | 
+142 141 | def bar5():
+
+RET505.py:145:5: RET505 [*] Unnecessary `else` after `return` statement
+    |
+143 |     if True:
+144 |         return
+145 |     else:  # comment
+    |     ^^^^ RET505
+146 |         pass
+    |
+    = help: Remove unnecessary `else`
+
+ℹ Safe fix
+142 142 | def bar5():
+143 143 |     if True:
+144 144 |         return
+145     |-    else:  # comment
+146     |-        pass
+    145 |+    # comment
+    146 |+    pass
+147 147 | 
+148 148 | 
+149 149 | def bar6():
+
+RET505.py:152:5: RET505 [*] Unnecessary `else` after `return` statement
+    |
+150 |     if True:
+151 |         return
+152 |     else\
+    |     ^^^^ RET505
+153 |         :\
+154 |         # comment
+    |
+    = help: Remove unnecessary `else`
+
+ℹ Safe fix
+149 149 | def bar6():
+150 150 |     if True:
+151 151 |         return
+152     |-    else\
+153     |-        :\
+154     |-        # comment
+155     |-        pass
+    152 |+    # comment
+    153 |+    pass
+156 154 | 
+157 155 | 
+158 156 | def bar7():
+
+RET505.py:161:5: RET505 [*] Unnecessary `else` after `return` statement
+    |
+159 |     if True:
+160 |         return
+161 |     else\
+    |     ^^^^ RET505
+162 |         :  # comment
+163 |         pass
+    |
+    = help: Remove unnecessary `else`
+
+ℹ Safe fix
+158 158 | def bar7():
+159 159 |     if True:
+160 160 |         return
+161     |-    else\
+162     |-        :  # comment
+163     |-        pass
+    161 |+    # comment
+    162 |+    pass
+164 163 | 
+165 164 | x = 0
+166 165 | 
+