[ty] Skip reading notebooks when discovering scripts (#28781)

crates/ruff_db/src/source.rs

@@ -22,7 +22,7 @@ pub fn source_text(db: &dyn Db, file: File) -> SourceText {
         return source.clone();
     }
 
-    let kind = if is_notebook(db.system(), path) {
+    let kind = if is_notebook_path(db.system(), path) {
         file.read_to_notebook(db)
             .unwrap_or_else(|error| {
                 tracing::debug!("Failed to read notebook '{path}': {error}");
@@ -47,7 +47,19 @@ pub fn source_text(db: &dyn Db, file: File) -> SourceText {
     }
 }
 
-fn is_notebook(system: &dyn System, path: &FilePath) -> bool {
+/// Returns whether a file is a notebook without reading its contents.
+pub fn is_notebook(db: &dyn Db, file: File) -> bool {
+    if let Some(source) = file.source_text_override(db) {
+        return source.is_notebook();
+    }
+
+    // The editor can change a file's source type without changing its path.
+    let _ = file.revision(db);
+
+    is_notebook_path(db.system(), file.path(db))
+}
+
+fn is_notebook_path(system: &dyn System, path: &FilePath) -> bool {
     let source_type = match path {
         FilePath::System(path) => system.source_type(path),
         FilePath::SystemVirtual(system_virtual) => system.virtual_path_source_type(system_virtual),

crates/ty_project/src/script.rs

@@ -4,7 +4,7 @@ use ruff_db::diagnostic::{
     Annotation, Diagnostic, DiagnosticId, Severity, Span, SubDiagnostic, SubDiagnosticSeverity,
 };
 use ruff_db::files::File;
-use ruff_db::source::source_text;
+use ruff_db::source::{is_notebook, source_text};
 use ruff_python_ast::script::ScriptTag;
 use ruff_ranged_value::{RangedValue, ValueSource, ValueSourceGuard};
 use ruff_text_size::{Ranged, TextRange, TextSize};
@@ -160,11 +160,12 @@ pub fn script_tag(db: &dyn SourceDb, file: File) -> Option<Box<ScriptTag>> {
         return None;
     }
 
-    let source = source_text(db, file);
-    if source.is_notebook() {
+    // Notebook outputs can be large, so skip notebooks before reading their contents.
+    if is_notebook(db, file) {
         return None;
     }
 
+    let source = source_text(db, file);
     ScriptTag::parse(source.as_bytes()).map(Box::new)
 }
 
@@ -394,8 +395,9 @@ fn invalid_script_metadata_diagnostic(
 #[cfg(test)]
 mod tests {
     use ruff_db::files::system_path_to_file;
+    use ruff_db::source::source_text;
     use ruff_db::system::{DbWithWritableSystem as _, SystemPath, SystemPathBuf};
-    use ruff_db::testing::assert_function_query_was_not_run;
+    use ruff_db::testing::{assert_function_query_was_not_run, assert_function_query_was_run};
     use ty_python_semantic::Db as _;
 
     use crate::db::testing::TestDb;
@@ -434,6 +436,30 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn script_tag_does_not_read_notebook_source() -> anyhow::Result<()> {
+        let mut db = TestDb::new(ProjectMetadata::new(
+            "test",
+            SystemPathBuf::from("/project"),
+        ));
+        db.write_file(
+            "/project/notebook.ipynb",
+            r#"{"cells": [], "metadata": {}, "nbformat": 4, "nbformat_minor": 5}"#,
+        )?;
+        let notebook = system_path_to_file(&db, SystemPath::new("/project/notebook.ipynb"))?;
+
+        assert_eq!(db.project().script_files(&db).iter().count(), 0);
+        assert!(db.project().files(&db).contains(notebook));
+
+        db.take_salsa_events();
+        source_text(&db, notebook);
+        let events = db.take_salsa_events();
+        // This assertion would fail if index already read the full text of notebook.
+        assert_function_query_was_run(&db, source_text, notebook, &events);
+
+        Ok(())
+    }
+
     #[test]
     fn equivalent_script_settings_share_programs() -> anyhow::Result<()> {
         let mut db = TestDb::new(ProjectMetadata::new(