mirror of
https://github.com/astral-sh/ruff.git
synced 2025-09-29 13:24:57 +00:00
Enable start-of-block insertions (#4741)
This commit is contained in:
Charlie Marsh
GitHub
parent
01470d9045
commit
bb4f3dedf4
1 changed files with 252 additions and 72 deletions
|
|
@ -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,56 +8,34 @@ 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) struct Insertion {
|
||||
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<'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.
|
||||
///
|
||||
/// For example, given the following code:
|
||||
///
|
||||
/// ```python
|
||||
/// """Hello, world!"""
|
||||
///
|
||||
/// import os
|
||||
/// import math
|
||||
///
|
||||
///
|
||||
/// def foo():
|
||||
/// pass
|
||||
/// ```
|
||||
///
|
||||
/// 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 {
|
||||
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(), ";")
|
||||
} else {
|
||||
// Otherwise, insert on the next line.
|
||||
Insertion::new(
|
||||
"",
|
||||
locator.full_line_end(location),
|
||||
stylist.line_ending().as_str(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// 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.
|
||||
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:
|
||||
///
|
||||
|
|
@ -65,17 +46,21 @@ impl Insertion {
|
|||
/// ```
|
||||
///
|
||||
/// 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 {
|
||||
/// 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;
|
||||
// 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(), ";");
|
||||
return Insertion::inline(" ", range.end(), ";");
|
||||
}
|
||||
|
||||
// Otherwise, advance to the next row.
|
||||
|
|
@ -95,25 +80,194 @@ impl Insertion {
|
|||
}
|
||||
}
|
||||
|
||||
Insertion::new("", location, stylist.line_ending().as_str())
|
||||
Insertion::own_line("", location, stylist.line_ending().as_str())
|
||||
}
|
||||
|
||||
fn new(prefix: &'static str, location: TextSize, suffix: &'static str) -> Self {
|
||||
Self {
|
||||
prefix,
|
||||
location,
|
||||
suffix,
|
||||
/// 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:
|
||||
///
|
||||
/// ```python
|
||||
/// """Hello, world!"""
|
||||
///
|
||||
/// import os
|
||||
/// import math
|
||||
///
|
||||
///
|
||||
/// def foo():
|
||||
/// pass
|
||||
/// ```
|
||||
///
|
||||
/// 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.
|
||||
///
|
||||
/// 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 statement is a semicolon, insert after the semicolon as
|
||||
// an inline statement.
|
||||
Insertion::inline(" ", range.end(), ";")
|
||||
} else {
|
||||
// Otherwise, insert on the next line.
|
||||
Insertion::own_line(
|
||||
"",
|
||||
locator.full_line_end(location),
|
||||
stylist.line_ending().as_str(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// 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
|
||||
/// if TYPE_CHECKING:
|
||||
/// import os
|
||||
/// ```
|
||||
///
|
||||
/// The insertion returned will begin at the start of the `import os` statement, and will
|
||||
/// 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,
|
||||
}
|
||||
|
||||
let mut state = Awaiting::Colon(0);
|
||||
for (tok, range) in
|
||||
lexer::lex_starts_at(locator.after(location), Mode::Module, location).flatten()
|
||||
{
|
||||
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.saturating_add(1));
|
||||
}
|
||||
Tok::Rpar | Tok::Rbrace | Tok::Rsqb => {
|
||||
state = Awaiting::Colon(depth.saturating_sub(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;
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// 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;
|
||||
Edit::insertion(format!("{prefix}{content}{suffix}"), location)
|
||||
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)
|
||||
}
|
||||
|
||||
/// 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,
|
||||
}
|
||||
}
|
||||
|
||||
/// 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,
|
||||
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),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -148,20 +302,20 @@ mod tests {
|
|||
|
||||
use super::Insertion;
|
||||
|
||||
fn insert(contents: &str) -> Result<Insertion> {
|
||||
let program = Suite::parse(contents, "<filename>")?;
|
||||
let tokens: Vec<LexResult> = 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<Insertion> {
|
||||
let program = Suite::parse(contents, "<filename>")?;
|
||||
let tokens: Vec<LexResult> = 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#"
|
||||
|
|
@ -169,7 +323,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#"
|
||||
|
|
@ -178,7 +332,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#"
|
||||
|
|
@ -188,7 +342,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#"
|
||||
|
|
@ -197,7 +351,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#"
|
||||
|
|
@ -206,7 +360,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#"
|
||||
|
|
@ -216,7 +370,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#"
|
||||
|
|
@ -226,7 +380,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#"
|
||||
|
|
@ -235,7 +389,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#"
|
||||
|
|
@ -244,7 +398,7 @@ x = 1
|
|||
.trim_start();
|
||||
assert_eq!(
|
||||
insert(contents)?,
|
||||
Insertion::new(" ", TextSize::from(20), ";")
|
||||
Insertion::inline(" ", TextSize::from(20), ";")
|
||||
);
|
||||
|
||||
let contents = r#"
|
||||
|
|
@ -254,9 +408,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<LexResult> = 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", " ")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue