From 9efd5b237e556f72e694b602f65937ce900b9c3c Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Tue, 30 May 2023 22:55:08 -0400 Subject: [PATCH] Enable start-of-block insertions --- crates/ruff/src/importer/insertion.rs | 312 ++++++++++++++++++++------ 1 file changed, 246 insertions(+), 66 deletions(-) diff --git a/crates/ruff/src/importer/insertion.rs b/crates/ruff/src/importer/insertion.rs index e2123ee16c0ce6..c66d663aade6a4 100644 --- a/crates/ruff/src/importer/insertion.rs +++ b/crates/ruff/src/importer/insertion.rs @@ -1,3 +1,6 @@ +//! Insert statements into Python code. +#![allow(dead_code)] + use ruff_text_size::TextSize; use rustpython_parser::ast::{Ranged, Stmt}; use rustpython_parser::{lexer, Mode, Tok}; @@ -5,20 +8,83 @@ use rustpython_parser::{lexer, Mode, Tok}; use ruff_diagnostics::Edit; use ruff_python_ast::helpers::is_docstring_stmt; use ruff_python_ast::source_code::{Locator, Stylist}; +use ruff_textwrap::indent; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) enum Placement<'a> { + /// The content will be inserted inline with the existing code (i.e., within semicolon-delimited + /// statements). + Inline, + /// The content will be inserted on its own line. + OwnLine, + /// The content will be inserted as an indented block. + Indented(&'a str), +} #[derive(Debug, Clone, PartialEq, Eq)] -pub(super) struct Insertion { +pub(super) struct Insertion<'a> { /// The content to add before the insertion. - prefix: &'static str, + prefix: &'a str, /// The location at which to insert. location: TextSize, /// The content to add after the insertion. - suffix: &'static str, + suffix: &'a str, + /// The line placement of insertion. + placement: Placement<'a>, } -impl Insertion { - /// Create an [`Insertion`] to insert (e.g.) an import after the end of the given [`Stmt`], - /// along with a prefix and suffix to use for the insertion. +impl<'a> Insertion<'a> { + /// Create an [`Insertion`] to insert (e.g.) an import statement at the start of a given + /// file, along with a prefix and suffix to use for the insertion. + /// + /// For example, given the following code: + /// + /// ```python + /// """Hello, world!""" + /// + /// import os + /// ``` + /// + /// The insertion returned will begin at the start of the `import os` statement, and will + /// include a trailing newline. + pub(super) fn start_of_file( + body: &[Stmt], + locator: &Locator, + stylist: &Stylist, + ) -> Insertion<'static> { + // Skip over any docstrings. + let mut location = if let Some(location) = match_docstring_end(body) { + // If the first token after the docstring is a semicolon, insert after the semicolon as an + // inline statement. + let first_token = lexer::lex_starts_at(locator.after(location), Mode::Module, location) + .flatten() + .next(); + if let Some((Tok::Semi, range)) = first_token { + return Insertion::inline(" ", range.end(), ";"); + } + + // Otherwise, advance to the next row. + locator.full_line_end(location) + } else { + TextSize::default() + }; + + // Skip over any comments and empty lines. + for (tok, range) in + lexer::lex_starts_at(locator.after(location), Mode::Module, location).flatten() + { + if matches!(tok, Tok::Comment(..) | Tok::Newline) { + location = locator.full_line_end(range.end()); + } else { + break; + } + } + + Insertion::own_line("", location, stylist.line_ending().as_str()) + } + + /// Create an [`Insertion`] to insert (e.g.) an import after the end of the given + /// [`Stmt`], along with a prefix and suffix to use for the insertion. /// /// For example, given the following code: /// @@ -34,18 +100,24 @@ impl Insertion { /// ``` /// /// The insertion returned will begin after the newline after the last import statement, which - /// in this case is the line after `import math`, and will include a trailing newline suffix. - pub(super) fn end_of_statement(stmt: &Stmt, locator: &Locator, stylist: &Stylist) -> Insertion { + /// in this case is the line after `import math`, and will include a trailing newline. + /// + /// The statement itself is assumed to be at the top-level of the module. + pub(super) fn end_of_statement( + stmt: &Stmt, + locator: &Locator, + stylist: &Stylist, + ) -> Insertion<'static> { let location = stmt.end(); let mut tokens = lexer::lex_starts_at(locator.after(location), Mode::Module, location).flatten(); if let Some((Tok::Semi, range)) = tokens.next() { - // If the first token after the docstring is a semicolon, insert after the semicolon as an - // inline statement; - Insertion::new(" ", range.end(), ";") + // If the first token after the statement is a semicolon, insert after the semicolon as + // an inline statement. + Insertion::inline(" ", range.end(), ";") } else { // Otherwise, insert on the next line. - Insertion::new( + Insertion::own_line( "", locator.full_line_end(location), stylist.line_ending().as_str(), @@ -53,67 +125,149 @@ impl Insertion { } } - /// Create an [`Insertion`] to insert (e.g.) an import statement at the "top" of a given file, - /// along with a prefix and suffix to use for the insertion. + /// Create an [`Insertion`] to insert (e.g.) an import statement at the start of a given + /// block, along with a prefix and suffix to use for the insertion. /// /// For example, given the following code: /// /// ```python - /// """Hello, world!""" - /// - /// import os + /// if TYPE_CHECKING: + /// import os /// ``` /// /// The insertion returned will begin at the start of the `import os` statement, and will - /// include a trailing newline suffix. - pub(super) fn start_of_file(body: &[Stmt], locator: &Locator, stylist: &Stylist) -> Insertion { - // Skip over any docstrings. - let mut location = if let Some(location) = match_docstring_end(body) { - // If the first token after the docstring is a semicolon, insert after the semicolon as an - // inline statement; - let first_token = lexer::lex_starts_at(locator.after(location), Mode::Module, location) - .flatten() - .next(); - if let Some((Tok::Semi, range)) = first_token { - return Insertion::new(" ", range.end(), ";"); - } - - // Otherwise, advance to the next row. - locator.full_line_end(location) - } else { - TextSize::default() - }; + /// include a trailing newline. + /// + /// The block itself is assumed to be at the top-level of the module. + pub(super) fn start_of_block( + mut location: TextSize, + locator: &Locator<'a>, + stylist: &Stylist, + ) -> Insertion<'a> { + enum Awaiting { + Colon(u32), + Newline, + Indent, + } - // Skip over any comments and empty lines. + let mut state = Awaiting::Colon(0); for (tok, range) in lexer::lex_starts_at(locator.after(location), Mode::Module, location).flatten() { - if matches!(tok, Tok::Comment(..) | Tok::Newline) { - location = locator.full_line_end(range.end()); - } else { - break; + match state { + // Iterate until we find the colon indicating the start of the block body. + Awaiting::Colon(depth) => match tok { + Tok::Colon if depth == 0 => { + state = Awaiting::Newline; + } + Tok::Lpar | Tok::Lbrace | Tok::Lsqb => { + state = Awaiting::Colon(depth + 1); + } + Tok::Rpar | Tok::Rbrace | Tok::Rsqb => { + state = Awaiting::Colon(depth - 1); + } + _ => {} + }, + // Once we've seen the colon, we're looking for a newline; otherwise, there's no + // block body (e.g. `if True: pass`). + Awaiting::Newline => match tok { + Tok::Comment(..) => {} + Tok::Newline => { + state = Awaiting::Indent; + } + _ => { + location = range.start(); + break; + } + }, + // Once we've seen the newline, we're looking for the indentation of the block body. + Awaiting::Indent => match tok { + Tok::NonLogicalNewline => {} + Tok::Indent => { + // This is like: + // ```py + // if True: + // pass + // ``` + // Where `range` is the indentation before the `pass` token. + return Insertion::indented( + "", + range.start(), + stylist.line_ending().as_str(), + locator.slice(range), + ); + } + _ => { + location = range.start(); + break; + } + }, } } - Insertion::new("", location, stylist.line_ending().as_str()) + // This is like: `if True: pass`, where `location` is the start of the `pass` token. + Insertion::inline("", location, "; ") + } + + /// Convert this [`Insertion`] into an [`Edit`] that inserts the given content. + pub(super) fn into_edit(self, content: &str) -> Edit { + let Insertion { + prefix, + location, + suffix, + placement, + } = self; + let content = format!("{prefix}{content}{suffix}"); + Edit::insertion( + match placement { + Placement::Indented(indentation) if !indentation.is_empty() => { + indent(&content, indentation).to_string() + } + _ => content, + }, + location, + ) + } + + /// Returns `true` if this [`Insertion`] is inline. + pub(super) fn is_inline(&self) -> bool { + matches!(self.placement, Placement::Inline) } - fn new(prefix: &'static str, location: TextSize, suffix: &'static str) -> Self { + /// Create an [`Insertion`] that inserts content inline (i.e., within semicolon-delimited + /// statements). + fn inline(prefix: &'a str, location: TextSize, suffix: &'a str) -> Self { Self { prefix, location, suffix, + placement: Placement::Inline, } } - /// Convert this [`Insertion`] into an [`Edit`] that inserts the given content. - pub(super) fn into_edit(self, content: &str) -> Edit { - let Insertion { + /// Create an [`Insertion`] that starts on its own line. + fn own_line(prefix: &'a str, location: TextSize, suffix: &'a str) -> Self { + Self { prefix, location, suffix, - } = self; - Edit::insertion(format!("{prefix}{content}{suffix}"), location) + placement: Placement::OwnLine, + } + } + + /// Create an [`Insertion`] that starts on its own line, with the given indentation. + fn indented( + prefix: &'a str, + location: TextSize, + suffix: &'a str, + indentation: &'a str, + ) -> Self { + Self { + prefix, + location, + suffix, + placement: Placement::Indented(indentation), + } } } @@ -147,20 +301,20 @@ mod tests { use super::Insertion; - fn insert(contents: &str) -> Result { - let program = parser::parse_program(contents, "")?; - let tokens: Vec = ruff_rustpython::tokenize(contents); - let locator = Locator::new(contents); - let stylist = Stylist::from_tokens(&tokens, &locator); - Ok(Insertion::start_of_file(&program, &locator, &stylist)) - } - #[test] fn start_of_file() -> Result<()> { + fn insert(contents: &str) -> Result { + let program = parser::parse_program(contents, "")?; + let tokens: Vec = ruff_rustpython::tokenize(contents); + let locator = Locator::new(contents); + let stylist = Stylist::from_tokens(&tokens, &locator); + Ok(Insertion::start_of_file(&program, &locator, &stylist)) + } + let contents = ""; assert_eq!( insert(contents)?, - Insertion::new("", TextSize::from(0), LineEnding::default().as_str()) + Insertion::own_line("", TextSize::from(0), LineEnding::default().as_str()) ); let contents = r#" @@ -168,7 +322,7 @@ mod tests { .trim_start(); assert_eq!( insert(contents)?, - Insertion::new("", TextSize::from(19), LineEnding::default().as_str()) + Insertion::own_line("", TextSize::from(19), LineEnding::default().as_str()) ); let contents = r#" @@ -177,7 +331,7 @@ mod tests { .trim_start(); assert_eq!( insert(contents)?, - Insertion::new("", TextSize::from(20), "\n") + Insertion::own_line("", TextSize::from(20), "\n") ); let contents = r#" @@ -187,7 +341,7 @@ mod tests { .trim_start(); assert_eq!( insert(contents)?, - Insertion::new("", TextSize::from(40), "\n") + Insertion::own_line("", TextSize::from(40), "\n") ); let contents = r#" @@ -196,7 +350,7 @@ x = 1 .trim_start(); assert_eq!( insert(contents)?, - Insertion::new("", TextSize::from(0), "\n") + Insertion::own_line("", TextSize::from(0), "\n") ); let contents = r#" @@ -205,7 +359,7 @@ x = 1 .trim_start(); assert_eq!( insert(contents)?, - Insertion::new("", TextSize::from(23), "\n") + Insertion::own_line("", TextSize::from(23), "\n") ); let contents = r#" @@ -215,7 +369,7 @@ x = 1 .trim_start(); assert_eq!( insert(contents)?, - Insertion::new("", TextSize::from(43), "\n") + Insertion::own_line("", TextSize::from(43), "\n") ); let contents = r#" @@ -225,7 +379,7 @@ x = 1 .trim_start(); assert_eq!( insert(contents)?, - Insertion::new("", TextSize::from(43), "\n") + Insertion::own_line("", TextSize::from(43), "\n") ); let contents = r#" @@ -234,7 +388,7 @@ x = 1 .trim_start(); assert_eq!( insert(contents)?, - Insertion::new("", TextSize::from(0), "\n") + Insertion::own_line("", TextSize::from(0), "\n") ); let contents = r#" @@ -243,7 +397,7 @@ x = 1 .trim_start(); assert_eq!( insert(contents)?, - Insertion::new(" ", TextSize::from(20), ";") + Insertion::inline(" ", TextSize::from(20), ";") ); let contents = r#" @@ -253,9 +407,35 @@ x = 1 .trim_start(); assert_eq!( insert(contents)?, - Insertion::new(" ", TextSize::from(20), ";") + Insertion::inline(" ", TextSize::from(20), ";") ); Ok(()) } + + #[test] + fn start_of_block() { + fn insert(contents: &str, offset: TextSize) -> Insertion { + let tokens: Vec = ruff_rustpython::tokenize(contents); + let locator = Locator::new(contents); + let stylist = Stylist::from_tokens(&tokens, &locator); + Insertion::start_of_block(offset, &locator, &stylist) + } + + let contents = "if True: pass"; + assert_eq!( + insert(contents, TextSize::from(0)), + Insertion::inline("", TextSize::from(9), "; ") + ); + + let contents = r#" +if True: + pass +"# + .trim_start(); + assert_eq!( + insert(contents, TextSize::from(0)), + Insertion::indented("", TextSize::from(9), "\n", " ") + ); + } }