diff --git a/crates/ruff_db/src/source.rs b/crates/ruff_db/src/source.rs index 511a7e7fa6..4849171765 100644 --- a/crates/ruff_db/src/source.rs +++ b/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), diff --git a/crates/ty_project/src/script.rs b/crates/ty_project/src/script.rs index 4c435ad905..1469f39d9e 100644 --- a/crates/ty_project/src/script.rs +++ b/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> { 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(