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

.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

crates/ty_server/tests/e2e/notebook.rs

@@ -1,5 +1,7 @@
 use insta::assert_json_snapshot;
-use lsp_types::{NotebookCellKind, Position, Range};
+use lsp_types::{CompletionResponse, CompletionTriggerKind, NotebookCellKind, Position, Range};
+use ruff_db::system::SystemPath;
+use ty_server::ClientOptions;
 
 use crate::{TestServer, TestServerBuilder};
 
@@ -276,6 +278,142 @@ fn swap_cells() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[test]
+fn auto_import() -> anyhow::Result<()> {
+    let mut server = TestServerBuilder::new()?
+        .with_workspace(
+            SystemPath::new("src"),
+            Some(ClientOptions::default().with_experimental_auto_import(true)),
+        )?
+        .build()?
+        .wait_until_workspaces_are_initialized()?;
+
+    server.initialization_result().unwrap();
+
+    let mut builder = NotebookBuilder::virtual_file("src/test.ipynb");
+
+    builder.add_python_cell(
+        r#"from typing import TYPE_CHECKING
+"#,
+    );
+
+    let second_cell = builder.add_python_cell(
+        r#"# leading comment
+b: Litera
+"#,
+    );
+
+    builder.open(&mut server);
+
+    server.collect_publish_diagnostic_notifications(2)?;
+
+    let completions = literal_completions(&mut server, &second_cell, Position::new(1, 9))?;
+
+    assert_json_snapshot!(completions);
+
+    Ok(())
+}
+
+#[test]
+fn auto_import_same_cell() -> anyhow::Result<()> {
+    let mut server = TestServerBuilder::new()?
+        .with_workspace(
+            SystemPath::new("src"),
+            Some(ClientOptions::default().with_experimental_auto_import(true)),
+        )?
+        .build()?
+        .wait_until_workspaces_are_initialized()?;
+
+    server.initialization_result().unwrap();
+
+    let mut builder = NotebookBuilder::virtual_file("src/test.ipynb");
+
+    let first_cell = builder.add_python_cell(
+        r#"from typing import TYPE_CHECKING
+b: Litera
+"#,
+    );
+
+    builder.open(&mut server);
+
+    server.collect_publish_diagnostic_notifications(1)?;
+
+    let completions = literal_completions(&mut server, &first_cell, Position::new(1, 9))?;
+
+    assert_json_snapshot!(completions);
+
+    Ok(())
+}
+
+#[test]
+fn auto_import_from_future() -> anyhow::Result<()> {
+    let mut server = TestServerBuilder::new()?
+        .with_workspace(
+            SystemPath::new("src"),
+            Some(ClientOptions::default().with_experimental_auto_import(true)),
+        )?
+        .build()?
+        .wait_until_workspaces_are_initialized()?;
+
+    server.initialization_result().unwrap();
+
+    let mut builder = NotebookBuilder::virtual_file("src/test.ipynb");
+
+    builder.add_python_cell(r#"from typing import TYPE_CHECKING"#);
+
+    let second_cell = builder.add_python_cell(
+        r#"from __future__ import annotations
+b: Litera
+"#,
+    );
+
+    builder.open(&mut server);
+
+    server.collect_publish_diagnostic_notifications(2)?;
+
+    let completions = literal_completions(&mut server, &second_cell, Position::new(1, 9))?;
+
+    assert_json_snapshot!(completions);
+
+    Ok(())
+}
+
+#[test]
+fn auto_import_docstring() -> anyhow::Result<()> {
+    let mut server = TestServerBuilder::new()?
+        .with_workspace(
+            SystemPath::new("src"),
+            Some(ClientOptions::default().with_experimental_auto_import(true)),
+        )?
+        .build()?
+        .wait_until_workspaces_are_initialized()?;
+
+    server.initialization_result().unwrap();
+
+    let mut builder = NotebookBuilder::virtual_file("src/test.ipynb");
+
+    builder.add_python_cell(
+        r#"from typing import TYPE_CHECKING
+"#,
+    );
+
+    let second_cell = builder.add_python_cell(
+        r#""""A cell level docstring"""
+b: Litera
+"#,
+    );
+
+    builder.open(&mut server);
+
+    server.collect_publish_diagnostic_notifications(2)?;
+
+    let completions = literal_completions(&mut server, &second_cell, Position::new(1, 9))?;
+
+    assert_json_snapshot!(completions);
+
+    Ok(())
+}
+
 fn semantic_tokens_full_for_cell(
     server: &mut TestServer,
     cell_uri: &lsp_types::Url,
@@ -359,3 +497,37 @@ impl NotebookBuilder {
         self.notebook_url
     }
 }
+
+fn literal_completions(
+    server: &mut TestServer,
+    cell: &lsp_types::Url,
+    position: Position,
+) -> crate::Result<Vec<lsp_types::CompletionItem>> {
+    let completions_id =
+        server.send_request::<lsp_types::request::Completion>(lsp_types::CompletionParams {
+            text_document_position: lsp_types::TextDocumentPositionParams {
+                text_document: lsp_types::TextDocumentIdentifier { uri: cell.clone() },
+                position,
+            },
+            work_done_progress_params: lsp_types::WorkDoneProgressParams::default(),
+            partial_result_params: lsp_types::PartialResultParams::default(),
+            context: Some(lsp_types::CompletionContext {
+                trigger_kind: CompletionTriggerKind::TRIGGER_FOR_INCOMPLETE_COMPLETIONS,
+                trigger_character: None,
+            }),
+        });
+
+    // There are a ton of imports we don't care about in here...
+    // The import bit is that an edit is always restricted to the current cell. That means,
+    // we can't add `Literal` to the `from typing import TYPE_CHECKING` import in cell 1
+    let completions = server.await_response::<lsp_types::request::Completion>(&completions_id)?;
+    let mut items = match completions {
+        Some(CompletionResponse::Array(array)) => array,
+        Some(CompletionResponse::List(lsp_types::CompletionList { items, .. })) => items,
+        None => return Ok(vec![]),
+    };
+
+    items.retain(|item| item.label.starts_with("Litera"));
+
+    Ok(items)
+}

crates/ty_server/tests/e2e/snapshots/e2e__notebook__auto_import.snap

@@ -0,0 +1,90 @@
+---
+source: crates/ty_server/tests/e2e/notebook.rs
+expression: completions
+---
+[
+  {
+    "label": "Literal (import typing)",
+    "kind": 6,
+    "sortText": " 43",
+    "insertText": "Literal",
+    "additionalTextEdits": [
+      {
+        "range": {
+          "start": {
+            "line": 1,
+            "character": 0
+          },
+          "end": {
+            "line": 1,
+            "character": 0
+          }
+        },
+        "newText": "from typing import Literal\n"
+      }
+    ]
+  },
+  {
+    "label": "Literal (import typing_extensions)",
+    "kind": 6,
+    "sortText": " 44",
+    "insertText": "Literal",
+    "additionalTextEdits": [
+      {
+        "range": {
+          "start": {
+            "line": 1,
+            "character": 0
+          },
+          "end": {
+            "line": 1,
+            "character": 0
+          }
+        },
+        "newText": "from typing_extensions import Literal\n"
+      }
+    ]
+  },
+  {
+    "label": "LiteralString (import typing)",
+    "kind": 6,
+    "sortText": " 45",
+    "insertText": "LiteralString",
+    "additionalTextEdits": [
+      {
+        "range": {
+          "start": {
+            "line": 1,
+            "character": 0
+          },
+          "end": {
+            "line": 1,
+            "character": 0
+          }
+        },
+        "newText": "from typing import LiteralString\n"
+      }
+    ]
+  },
+  {
+    "label": "LiteralString (import typing_extensions)",
+    "kind": 6,
+    "sortText": " 46",
+    "insertText": "LiteralString",
+    "additionalTextEdits": [
+      {
+        "range": {
+          "start": {
+            "line": 1,
+            "character": 0
+          },
+          "end": {
+            "line": 1,
+            "character": 0
+          }
+        },
+        "newText": "from typing_extensions import LiteralString\n"
+      }
+    ]
+  }
+]

crates/ty_server/tests/e2e/snapshots/e2e__notebook__auto_import_docstring.snap

@@ -0,0 +1,90 @@
+---
+source: crates/ty_server/tests/e2e/notebook.rs
+expression: completions
+---
+[
+  {
+    "label": "Literal (import typing)",
+    "kind": 6,
+    "sortText": " 43",
+    "insertText": "Literal",
+    "additionalTextEdits": [
+      {
+        "range": {
+          "start": {
+            "line": 1,
+            "character": 0
+          },
+          "end": {
+            "line": 1,
+            "character": 0
+          }
+        },
+        "newText": "from typing import Literal\n"
+      }
+    ]
+  },
+  {
+    "label": "Literal (import typing_extensions)",
+    "kind": 6,
+    "sortText": " 44",
+    "insertText": "Literal",
+    "additionalTextEdits": [
+      {
+        "range": {
+          "start": {
+            "line": 1,
+            "character": 0
+          },
+          "end": {
+            "line": 1,
+            "character": 0
+          }
+        },
+        "newText": "from typing_extensions import Literal\n"
+      }
+    ]
+  },
+  {
+    "label": "LiteralString (import typing)",
+    "kind": 6,
+    "sortText": " 45",
+    "insertText": "LiteralString",
+    "additionalTextEdits": [
+      {
+        "range": {
+          "start": {
+            "line": 1,
+            "character": 0
+          },
+          "end": {
+            "line": 1,
+            "character": 0
+          }
+        },
+        "newText": "from typing import LiteralString\n"
+      }
+    ]
+  },
+  {
+    "label": "LiteralString (import typing_extensions)",
+    "kind": 6,
+    "sortText": " 46",
+    "insertText": "LiteralString",
+    "additionalTextEdits": [
+      {
+        "range": {
+          "start": {
+            "line": 1,
+            "character": 0
+          },
+          "end": {
+            "line": 1,
+            "character": 0
+          }
+        },
+        "newText": "from typing_extensions import LiteralString\n"
+      }
+    ]
+  }
+]

crates/ty_server/tests/e2e/snapshots/e2e__notebook__auto_import_from_future.snap

@@ -0,0 +1,90 @@
+---
+source: crates/ty_server/tests/e2e/notebook.rs
+expression: completions
+---
+[
+  {
+    "label": "Literal (import typing)",
+    "kind": 6,
+    "sortText": " 43",
+    "insertText": "Literal",
+    "additionalTextEdits": [
+      {
+        "range": {
+          "start": {
+            "line": 1,
+            "character": 0
+          },
+          "end": {
+            "line": 1,
+            "character": 0
+          }
+        },
+        "newText": "from typing import Literal\n"
+      }
+    ]
+  },
+  {
+    "label": "Literal (import typing_extensions)",
+    "kind": 6,
+    "sortText": " 44",
+    "insertText": "Literal",
+    "additionalTextEdits": [
+      {
+        "range": {
+          "start": {
+            "line": 1,
+            "character": 0
+          },
+          "end": {
+            "line": 1,
+            "character": 0
+          }
+        },
+        "newText": "from typing_extensions import Literal\n"
+      }
+    ]
+  },
+  {
+    "label": "LiteralString (import typing)",
+    "kind": 6,
+    "sortText": " 45",
+    "insertText": "LiteralString",
+    "additionalTextEdits": [
+      {
+        "range": {
+          "start": {
+            "line": 1,
+            "character": 0
+          },
+          "end": {
+            "line": 1,
+            "character": 0
+          }
+        },
+        "newText": "from typing import LiteralString\n"
+      }
+    ]
+  },
+  {
+    "label": "LiteralString (import typing_extensions)",
+    "kind": 6,
+    "sortText": " 46",
+    "insertText": "LiteralString",
+    "additionalTextEdits": [
+      {
+        "range": {
+          "start": {
+            "line": 1,
+            "character": 0
+          },
+          "end": {
+            "line": 1,
+            "character": 0
+          }
+        },
+        "newText": "from typing_extensions import LiteralString\n"
+      }
+    ]
+  }
+]

crates/ty_server/tests/e2e/snapshots/e2e__notebook__auto_import_same_cell.snap

@@ -0,0 +1,90 @@
+---
+source: crates/ty_server/tests/e2e/notebook.rs
+expression: completions
+---
+[
+  {
+    "label": "Literal (import typing)",
+    "kind": 6,
+    "sortText": " 43",
+    "insertText": "Literal",
+    "additionalTextEdits": [
+      {
+        "range": {
+          "start": {
+            "line": 0,
+            "character": 32
+          },
+          "end": {
+            "line": 0,
+            "character": 32
+          }
+        },
+        "newText": ", Literal"
+      }
+    ]
+  },
+  {
+    "label": "Literal (import typing_extensions)",
+    "kind": 6,
+    "sortText": " 44",
+    "insertText": "Literal",
+    "additionalTextEdits": [
+      {
+        "range": {
+          "start": {
+            "line": 0,
+            "character": 0
+          },
+          "end": {
+            "line": 0,
+            "character": 0
+          }
+        },
+        "newText": "from typing_extensions import Literal\n"
+      }
+    ]
+  },
+  {
+    "label": "LiteralString (import typing)",
+    "kind": 6,
+    "sortText": " 45",
+    "insertText": "LiteralString",
+    "additionalTextEdits": [
+      {
+        "range": {
+          "start": {
+            "line": 0,
+            "character": 32
+          },
+          "end": {
+            "line": 0,
+            "character": 32
+          }
+        },
+        "newText": ", LiteralString"
+      }
+    ]
+  },
+  {
+    "label": "LiteralString (import typing_extensions)",
+    "kind": 6,
+    "sortText": " 46",
+    "insertText": "LiteralString",
+    "additionalTextEdits": [
+      {
+        "range": {
+          "start": {
+            "line": 0,
+            "character": 0
+          },
+          "end": {
+            "line": 0,
+            "character": 0
+          }
+        },
+        "newText": "from typing_extensions import LiteralString\n"
+      }
+    ]
+  }
+]