Enable start-of-block insertions (#4741)

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,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 {
+impl<'a> Insertion<'a> {
-    /// Create an [`Insertion`] to insert (e.g.) an import after the end of the given [`Stmt`],
+    /// Create an [`Insertion`] to insert (e.g.) an import statement at the start of a given
-    /// along with a prefix and suffix to use for the insertion.
+    /// file, 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.
     ///
     /// 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.
+    /// include a trailing newline.
-    pub(super) fn start_of_file(body: &[Stmt], locator: &Locator, stylist: &Stylist) -> Insertion {
+    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 {
+    /// Create an [`Insertion`] to insert (e.g.) an import after the end of the given
-        Self {
+    /// [`Stmt`], along with a prefix and suffix to use for the insertion.
-            prefix,
+    ///
-            location,
+    /// For example, given the following code:
-            suffix,
+    ///
+    /// ```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", "    ")
+        );
+    }
 }