From eb42ce93198a1b7fa9f248feee9a8d75ca1b4ca9 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Fri, 3 Mar 2023 14:15:55 -0500 Subject: [PATCH] Extend `RET503` autofixes to "end of statement", including comments (#3324) --- .../test/fixtures/flake8_return/RET503.py | 27 +++++ crates/ruff/src/ast/helpers.rs | 26 +++++ crates/ruff/src/rules/flake8_return/rules.rs | 19 ++- ...lake8_return__tests__RET503_RET503.py.snap | 109 ++++++++++++++++-- crates/ruff_dev/src/print_tokens.rs | 4 +- 5 files changed, 167 insertions(+), 18 deletions(-) diff --git a/crates/ruff/resources/test/fixtures/flake8_return/RET503.py b/crates/ruff/resources/test/fixtures/flake8_return/RET503.py index 49a2c2213c..77d59cd8d4 100644 --- a/crates/ruff/resources/test/fixtures/flake8_return/RET503.py +++ b/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 diff --git a/crates/ruff/src/ast/helpers.rs b/crates/ruff/src/ast/helpers.rs index 22001d310a..387e91c45f 100644 --- a/crates/ruff/src/ast/helpers.rs +++ b/crates/ruff/src/ast/helpers.rs @@ -1113,6 +1113,32 @@ pub fn first_colon_range(range: Range, locator: &Locator) -> Option { 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 { let StmtKind::If { body, orelse, .. } = &stmt.node else { diff --git a/crates/ruff/src/rules/flake8_return/rules.rs b/crates/ruff/src/rules/flake8_return/rules.rs index f8f880d995..8ab2da5dcd 100644 --- a/crates/ruff/src/rules/flake8_return/rules.rs +++ b/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); diff --git a/crates/ruff/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET503_RET503.py.snap b/crates/ruff/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET503_RET503.py.snap index 71b17d7822..c53ef9a949 100644 --- a/crates/ruff/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET503_RET503.py.snap +++ b/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: ~ diff --git a/crates/ruff_dev/src/print_tokens.rs b/crates/ruff_dev/src/print_tokens.rs index af41f32b62..5e771e746f 100644 --- a/crates/ruff_dev/src/print_tokens.rs +++ b/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(()) }