[ty] Respect notebook cell boundaries when adding an auto import (#21322)

crates/ruff_python_importer/src/insertion.rs

@@ -10,7 +10,7 @@ use ruff_python_parser::{TokenKind, Tokens};
 use ruff_python_trivia::is_python_whitespace;
 use ruff_python_trivia::{PythonWhitespace, textwrap::indent};
 use ruff_source_file::{LineRanges, UniversalNewlineIterator};
-use ruff_text_size::{Ranged, TextSize};
+use ruff_text_size::{Ranged, TextRange, TextSize};
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(super) enum Placement<'a> {
@@ -37,7 +37,7 @@ pub struct Insertion<'a> {
 
 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.
+    /// file or cell, along with a prefix and suffix to use for the insertion.
     ///
     /// For example, given the following code:
     ///
@@ -49,7 +49,26 @@ impl<'a> Insertion<'a> {
     ///
     /// The insertion returned will begin at the start of the `import os` statement, and will
     /// include a trailing newline.
-    pub fn start_of_file(body: &[Stmt], contents: &str, stylist: &Stylist) -> Insertion<'static> {
+    ///
+    /// If `within_range` is set, the insertion will be limited to the specified range. That is,
+    /// the insertion is constrained to the given range rather than the start of the file.
+    /// This is used for insertions in notebook cells where the source code and AST are for
+    /// the entire notebook but the insertion should be constrained to a specific cell.
+    pub fn start_of_file(
+        body: &[Stmt],
+        contents: &str,
+        stylist: &Stylist,
+        within_range: Option<TextRange>,
+    ) -> Insertion<'static> {
+        let body = within_range
+            .map(|range| {
+                let start = body.partition_point(|stmt| stmt.start() < range.start());
+                let end = body.partition_point(|stmt| stmt.end() <= range.end());
+
+                &body[start..end]
+            })
+            .unwrap_or(body);
+
         // Skip over any docstrings.
         let mut location = if let Some(mut location) = match_docstring_end(body) {
             // If the first token after the docstring is a semicolon, insert after the semicolon as
@@ -66,6 +85,10 @@ impl<'a> Insertion<'a> {
 
             // Otherwise, advance to the next row.
             contents.full_line_end(location)
+        } else if let Some(range) = within_range
+            && range.start() != TextSize::ZERO
+        {
+            range.start()
         } else {
             contents.bom_start_offset()
         };
@@ -374,7 +397,12 @@ mod tests {
         fn insert(contents: &str) -> Result<Insertion<'_>> {
             let parsed = parse_module(contents)?;
             let stylist = Stylist::from_tokens(parsed.tokens(), contents);
-            Ok(Insertion::start_of_file(parsed.suite(), contents, &stylist))
+            Ok(Insertion::start_of_file(
+                parsed.suite(),
+                contents,
+                &stylist,
+                None,
+            ))
         }
 
         let contents = "";