Extend RET503 autofixes to "end of statement", including comments (#3324)

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

@@ -293,3 +293,30 @@ def x(y):
 
 def foo(baz: str) -> str:
     return baz
+
+
+def end_of_statement():
+    def example():
+        if True:
+            return ""
+
+
+    def example():
+        if True:
+            return ""
+
+
+    def example():
+        if True:
+            return ""  # type: ignore
+
+
+    def example():
+        if True:
+            return ""  ;
+
+
+    def example():
+        if True:
+            return "" \
+                ;  # type: ignore

crates/ruff/src/ast/helpers.rs

@@ -1113,6 +1113,32 @@ pub fn first_colon_range(range: Range, locator: &Locator) -> Option<Range> {
     range
 }
 
+/// Given a statement, find its "logical end".
+///
+/// For example: the statement could be following by a trailing semicolon, by an end-of-line
+/// comment, or by any number of continuation lines (and then by a comment, and so on).
+pub fn end_of_statement(stmt: &Stmt, locator: &Locator) -> Location {
+    let contents = locator.skip(stmt.end_location.unwrap());
+
+    // End-of-file, so just return the end of the statement.
+    if contents.is_empty() {
+        return stmt.end_location.unwrap();
+    }
+
+    // Otherwise, find the end of the last line that's "part of" the statement.
+    for (lineno, line) in contents.lines().enumerate() {
+        if line.ends_with('\\') {
+            continue;
+        }
+        return to_absolute(
+            Location::new(lineno + 1, line.chars().count()),
+            stmt.end_location.unwrap(),
+        );
+    }
+
+    unreachable!("Expected to find end-of-statement")
+}
+
 /// Return the `Range` of the first `Elif` or `Else` token in an `If` statement.
 pub fn elif_else_range(stmt: &Stmt, locator: &Locator) -> Option<Range> {
     let StmtKind::If { body, orelse, .. } = &stmt.node else {

crates/ruff/src/rules/flake8_return/rules.rs

@@ -3,7 +3,7 @@ use rustpython_parser::ast::{Constant, Expr, ExprKind, Location, Stmt, StmtKind}
 
 use ruff_macros::{define_violation, derive_message_formats};
 
-use crate::ast::helpers::elif_else_range;
+use crate::ast::helpers::{elif_else_range, end_of_statement};
 use crate::ast::types::Range;
 use crate::ast::visitor::Visitor;
 use crate::ast::whitespace::indentation;
@@ -216,7 +216,12 @@ fn implicit_return(checker: &mut Checker, stmt: &Stmt) {
                         content.push_str(checker.stylist.line_ending().as_str());
                         content.push_str(indent);
                         content.push_str("return None");
-                        diagnostic.amend(Fix::insertion(content, stmt.end_location.unwrap()));
+                        // This is the last statement in the function. So it has to be followed by
+                        // a newline, or comments, or nothing.
+                        diagnostic.amend(Fix::insertion(
+                            content,
+                            end_of_statement(stmt, checker.locator),
+                        ));
                     }
                 }
                 checker.diagnostics.push(diagnostic);
@@ -251,7 +256,10 @@ fn implicit_return(checker: &mut Checker, stmt: &Stmt) {
                         content.push_str(checker.stylist.line_ending().as_str());
                         content.push_str(indent);
                         content.push_str("return None");
-                        diagnostic.amend(Fix::insertion(content, stmt.end_location.unwrap()));
+                        diagnostic.amend(Fix::insertion(
+                            content,
+                            end_of_statement(stmt, checker.locator),
+                        ));
                     }
                 }
                 checker.diagnostics.push(diagnostic);
@@ -287,7 +295,10 @@ fn implicit_return(checker: &mut Checker, stmt: &Stmt) {
                     content.push_str(checker.stylist.line_ending().as_str());
                     content.push_str(indent);
                     content.push_str("return None");
-                    diagnostic.amend(Fix::insertion(content, stmt.end_location.unwrap()));
+                    diagnostic.amend(Fix::insertion(
+                        content,
+                        end_of_statement(stmt, checker.locator),
+                    ));
                 }
             }
             checker.diagnostics.push(diagnostic);

crates/ruff/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET503_RET503.py.snap

@@ -31,10 +31,10 @@ expression: diagnostics
     content: "\n        return None"
     location:
       row: 27
-      column: 15
+      column: 24
     end_location:
       row: 27
-      column: 15
+      column: 24
   parent: ~
 - kind:
     ImplicitReturn: ~
@@ -48,10 +48,10 @@ expression: diagnostics
     content: "\n    return None"
     location:
       row: 36
-      column: 11
+      column: 20
     end_location:
       row: 36
-      column: 11
+      column: 20
   parent: ~
 - kind:
     ImplicitReturn: ~
@@ -82,10 +82,10 @@ expression: diagnostics
     content: "\n        return None"
     location:
       row: 52
-      column: 15
+      column: 24
     end_location:
       row: 52
-      column: 15
+      column: 24
   parent: ~
 - kind:
     ImplicitReturn: ~
@@ -99,10 +99,10 @@ expression: diagnostics
     content: "\n    return None"
     location:
       row: 59
-      column: 22
+      column: 31
     end_location:
       row: 59
-      column: 22
+      column: 31
   parent: ~
 - kind:
     ImplicitReturn: ~
@@ -116,10 +116,10 @@ expression: diagnostics
     content: "\n    return None"
     location:
       row: 66
-      column: 21
+      column: 30
     end_location:
       row: 66
-      column: 21
+      column: 30
   parent: ~
 - kind:
     ImplicitReturn: ~
@@ -235,9 +235,94 @@ expression: diagnostics
     content: "\n            return None"
     location:
       row: 291
-      column: 19
+      column: 28
     end_location:
       row: 291
-      column: 19
+      column: 28
+  parent: ~
+- kind:
+    ImplicitReturn: ~
+  location:
+    row: 300
+    column: 8
+  end_location:
+    row: 301
+    column: 21
+  fix:
+    content: "\n        return None"
+    location:
+      row: 301
+      column: 21
+    end_location:
+      row: 301
+      column: 21
+  parent: ~
+- kind:
+    ImplicitReturn: ~
+  location:
+    row: 305
+    column: 8
+  end_location:
+    row: 306
+    column: 21
+  fix:
+    content: "\n        return None"
+    location:
+      row: 306
+      column: 21
+    end_location:
+      row: 306
+      column: 21
+  parent: ~
+- kind:
+    ImplicitReturn: ~
+  location:
+    row: 310
+    column: 8
+  end_location:
+    row: 311
+    column: 21
+  fix:
+    content: "\n        return None"
+    location:
+      row: 311
+      column: 37
+    end_location:
+      row: 311
+      column: 37
+  parent: ~
+- kind:
+    ImplicitReturn: ~
+  location:
+    row: 315
+    column: 8
+  end_location:
+    row: 316
+    column: 21
+  fix:
+    content: "\n        return None"
+    location:
+      row: 316
+      column: 24
+    end_location:
+      row: 316
+      column: 24
+  parent: ~
+- kind:
+    ImplicitReturn: ~
+  location:
+    row: 320
+    column: 8
+  end_location:
+    row: 321
+    column: 21
+  fix:
+    content: "\n        return None"
+    location:
+      row: 322
+      column: 33
+    end_location:
+      row: 322
+      column: 33
   parent: ~
 

crates/ruff_dev/src/print_tokens.rs

@@ -16,8 +16,8 @@ pub struct Args {
 
 pub fn main(args: &Args) -> Result<()> {
     let contents = fs::read_to_string(&args.file)?;
-    for (_, tok, _) in lexer::lex(&contents, Mode::Module).flatten() {
-        println!("{tok:#?}");
+    for (start, tok, end) in lexer::lex(&contents, Mode::Module).flatten() {
+        println!("{start:#?} {tok:#?} {end:#?}");
     }
     Ok(())
 }