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

Author: MichaReiser

.config/nextest.toml

@@ -7,6 +7,10 @@ serial = { max-threads = 1 }
 filter = 'binary(file_watching)'
 test-group = 'serial'
 
+[[profile.default.overrides]]
+filter = 'binary(e2e)'
+test-group = 'serial'
+
 [profile.ci]
 # Print out output for failing tests as soon as they fail, and also at the end
 # of the run (for easy scrollability).

crates/ruff_linter/src/importer/mod.rs

@@ -83,7 +83,7 @@ impl<'a> Importer<'a> {
                     .into_edit(&required_import)
             } else {
                 // Insert at the start of the file.
-                Insertion::start_of_file(self.python_ast, self.source, self.stylist)
+                Insertion::start_of_file(self.python_ast, self.source, self.stylist, None)
                     .into_edit(&required_import)
             }
         }
@@ -113,7 +113,7 @@ impl<'a> Importer<'a> {
             Insertion::end_of_statement(stmt, self.source, self.stylist)
         } else {
             // Insert at the start of the file.
-            Insertion::start_of_file(self.python_ast, self.source, self.stylist)
+            Insertion::start_of_file(self.python_ast, self.source, self.stylist, None)
         };
         let add_import_edit = insertion.into_edit(&content);
 
@@ -498,7 +498,7 @@ impl<'a> Importer<'a> {
             Insertion::end_of_statement(stmt, self.source, self.stylist)
         } else {
             // Insert at the start of the file.
-            Insertion::start_of_file(self.python_ast, self.source, self.stylist)
+            Insertion::start_of_file(self.python_ast, self.source, self.stylist, None)
         };
         if insertion.is_inline() {
             Err(anyhow::anyhow!(

crates/ruff_notebook/src/cell.rs

@@ -294,19 +294,33 @@ impl CellOffsets {
     }
 
     /// Returns `true` if the given range contains a cell boundary.
+    ///
+    /// A range starting at the cell boundary isn't considered to contain the cell boundary
+    /// as it starts right after it. A range starting before a cell boundary
+    /// and ending exactly at the boundary is considered to contain the cell boundary.
+    ///
+    /// # Examples
+    /// Cell 1:
+    ///
+    /// ```py
+    /// import c
+    /// ```
+    ///
+    /// Cell 2:
+    ///
+    /// ```py
+    /// import os
+    /// ```
+    ///
+    /// The range `import c`..`import os`, contains a cell boundary because it starts before cell 2 and ends in cell 2 (`end=cell2_boundary`).
+    /// The `import os` contains no cell boundary because it starts at the start of cell 2 (at the cell boundary) but doesn't cross into another cell.
     pub fn has_cell_boundary(&self, range: TextRange) -> bool {
-        self.binary_search_by(|offset| {
-            if range.start() <= *offset {
-                if range.end() < *offset {
-                    std::cmp::Ordering::Greater
-                } else {
-                    std::cmp::Ordering::Equal
-                }
-            } else {
-                std::cmp::Ordering::Less
-            }
-        })
-        .is_ok()
+        let after_range_start = self.partition_point(|offset| *offset <= range.start());
+        let Some(boundary) = self.get(after_range_start).copied() else {
+            return false;
+        };
+
+        range.contains_inclusive(boundary)
     }
 
     /// Returns an iterator over [`TextRange`]s covered by each cell.

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 = "";

crates/ty_ide/src/importer.rs

@@ -20,6 +20,7 @@ use rustc_hash::FxHashMap;
 
 use ruff_db::files::File;
 use ruff_db::parsed::ParsedModuleRef;
+use ruff_db::source::source_text;
 use ruff_diagnostics::Edit;
 use ruff_python_ast as ast;
 use ruff_python_ast::name::Name;
@@ -76,6 +77,7 @@ impl<'a> Importer<'a> {
         parsed: &'a ParsedModuleRef,
     ) -> Self {
         let imports = TopLevelImports::find(parsed);
+
         Self {
             db,
             file,
@@ -145,10 +147,14 @@ impl<'a> Importer<'a> {
         let request = request.avoid_conflicts(self.db, self.file, members);
         let mut symbol_text: Box<str> = request.member.into();
         let Some(response) = self.find(&request, members.at) else {
-            let insertion = if let Some(future) = self.find_last_future_import() {
+            let insertion = if let Some(future) = self.find_last_future_import(members.at) {
                 Insertion::end_of_statement(future.stmt, self.source, self.stylist)
             } else {
-                Insertion::start_of_file(self.parsed.suite(), self.source, self.stylist)
+                let range = source_text(self.db, self.file)
+                    .as_notebook()
+                    .and_then(|notebook| notebook.cell_offsets().containing_range(members.at));
+
+                Insertion::start_of_file(self.parsed.suite(), self.source, self.stylist, range)
             };
             let import = insertion.into_edit(&request.to_string());
             if matches!(request.style, ImportStyle::Import) {
@@ -209,6 +215,9 @@ impl<'a> Importer<'a> {
         available_at: TextSize,
     ) -> Option<ImportResponse<'importer, 'a>> {
         let mut choice = None;
+        let source = source_text(self.db, self.file);
+        let notebook = source.as_notebook();
+
         for import in &self.imports {
             // If the import statement comes after the spot where we
             // need the symbol, then we conservatively assume that
@@ -226,7 +235,22 @@ impl<'a> Importer<'a> {
             if import.stmt.start() >= available_at {
                 return choice;
             }
+
             if let Some(response) = import.satisfies(self.db, self.file, request) {
+                let partial = matches!(response.kind, ImportResponseKind::Partial { .. });
+
+                // The LSP doesn't support edits across cell boundaries.
+                // Skip over imports that only partially satisfy the import
+                // because they would require changes to the import (across cell boundaries).
+                if partial
+                    && let Some(notebook) = notebook
+                    && notebook
+                        .cell_offsets()
+                        .has_cell_boundary(TextRange::new(import.stmt.start(), available_at))
+                {
+                    continue;
+                }
+
                 if choice
                     .as_ref()
                     .is_none_or(|c| !c.kind.is_prioritized_over(&response.kind))
@@ -247,9 +271,21 @@ impl<'a> Importer<'a> {
     }
 
     /// Find the last `from __future__` import statement in the AST.
-    fn find_last_future_import(&self) -> Option<&'a AstImport> {
+    fn find_last_future_import(&self, at: TextSize) -> Option<&'a AstImport> {
+        let source = source_text(self.db, self.file);
+        let notebook = source.as_notebook();
+
         self.imports
             .iter()
+            .take_while(|import| import.stmt.start() <= at)
+            // Skip over imports from other cells.
+            .skip_while(|import| {
+                notebook.is_some_and(|notebook| {
+                    notebook
+                        .cell_offsets()
+                        .has_cell_boundary(TextRange::new(import.stmt.start(), at))
+                })
+            })
             .take_while(|import| {
                 import
                     .stmt