Refactor SourceKind to store file content (#6640)

crates/ruff/src/jupyter/notebook.rs

@@ -424,12 +424,13 @@ impl Notebook {
     }
 
     /// Update the notebook with the given sourcemap and transformed content.
-    pub(crate) fn update(&mut self, source_map: &SourceMap, transformed: &str) {
+    pub(crate) fn update(&mut self, source_map: &SourceMap, transformed: String) {
         // Cell offsets must be updated before updating the cell content as
         // it depends on the offsets to extract the cell content.
+        self.index.take();
         self.update_cell_offsets(source_map);
-        self.update_cell_content(transformed);
-        self.source_code = transformed.to_string();
+        self.update_cell_content(&transformed);
+        self.source_code = transformed;
     }
 
     /// Return a slice of [`Cell`] in the Jupyter notebook.
@@ -476,7 +477,9 @@ mod tests {
     use crate::jupyter::schema::Cell;
     use crate::jupyter::Notebook;
     use crate::registry::Rule;
-    use crate::test::{read_jupyter_notebook, test_notebook_path, test_resource_path};
+    use crate::test::{
+        read_jupyter_notebook, test_notebook_path, test_resource_path, TestedNotebook,
+    };
     use crate::{assert_messages, settings};
 
     /// Read a Jupyter cell from the `resources/test/fixtures/jupyter/cell` directory.
@@ -578,49 +581,64 @@ print("after empty cells")
     #[test]
     fn test_import_sorting() -> Result<()> {
         let path = "isort.ipynb".to_string();
-        let (diagnostics, source_kind, _) = test_notebook_path(
+        let TestedNotebook {
+            messages,
+            source_notebook,
+            ..
+        } = test_notebook_path(
             &path,
             Path::new("isort_expected.ipynb"),
             &settings::Settings::for_rule(Rule::UnsortedImports),
         )?;
-        assert_messages!(diagnostics, path, source_kind);
+        assert_messages!(messages, path, source_notebook);
         Ok(())
     }
 
     #[test]
     fn test_ipy_escape_command() -> Result<()> {
         let path = "ipy_escape_command.ipynb".to_string();
-        let (diagnostics, source_kind, _) = test_notebook_path(
+        let TestedNotebook {
+            messages,
+            source_notebook,
+            ..
+        } = test_notebook_path(
             &path,
             Path::new("ipy_escape_command_expected.ipynb"),
             &settings::Settings::for_rule(Rule::UnusedImport),
         )?;
-        assert_messages!(diagnostics, path, source_kind);
+        assert_messages!(messages, path, source_notebook);
         Ok(())
     }
 
     #[test]
     fn test_unused_variable() -> Result<()> {
         let path = "unused_variable.ipynb".to_string();
-        let (diagnostics, source_kind, _) = test_notebook_path(
+        let TestedNotebook {
+            messages,
+            source_notebook,
+            ..
+        } = test_notebook_path(
             &path,
             Path::new("unused_variable_expected.ipynb"),
             &settings::Settings::for_rule(Rule::UnusedVariable),
         )?;
-        assert_messages!(diagnostics, path, source_kind);
+        assert_messages!(messages, path, source_notebook);
         Ok(())
     }
 
     #[test]
     fn test_json_consistency() -> Result<()> {
         let path = "before_fix.ipynb".to_string();
-        let (_, _, source_kind) = test_notebook_path(
+        let TestedNotebook {
+            linted_notebook: fixed_notebook,
+            ..
+        } = test_notebook_path(
             path,
             Path::new("after_fix.ipynb"),
             &settings::Settings::for_rule(Rule::UnusedImport),
         )?;
         let mut writer = Vec::new();
-        source_kind.expect_jupyter().write_inner(&mut writer)?;
+        fixed_notebook.write_inner(&mut writer)?;
         let actual = String::from_utf8(writer)?;
         let expected =
             std::fs::read_to_string(test_resource_path("fixtures/jupyter/after_fix.ipynb"))?;

crates/ruff/src/linter.rs

@@ -63,7 +63,7 @@ pub struct FixerResult<'a> {
     /// The result returned by the linter, after applying any fixes.
     pub result: LinterResult<(Vec<Message>, Option<ImportMap>)>,
     /// The resulting source code, after applying any fixes.
-    pub transformed: Cow<'a, str>,
+    pub transformed: Cow<'a, SourceKind>,
     /// The number of fixes applied for each [`Rule`].
     pub fixed: FixTable,
 }
@@ -335,19 +335,19 @@ pub fn add_noqa_to_path(path: &Path, package: Option<&Path>, settings: &Settings
 /// Generate a [`Message`] for each [`Diagnostic`] triggered by the given source
 /// code.
 pub fn lint_only(
-    contents: &str,
     path: &Path,
     package: Option<&Path>,
     settings: &Settings,
     noqa: flags::Noqa,
-    source_kind: Option<&SourceKind>,
+    source_kind: &SourceKind,
     source_type: PySourceType,
 ) -> LinterResult<(Vec<Message>, Option<ImportMap>)> {
     // Tokenize once.
-    let tokens: Vec<LexResult> = ruff_python_parser::tokenize(contents, source_type.as_mode());
+    let tokens: Vec<LexResult> =
+        ruff_python_parser::tokenize(source_kind.source_code(), source_type.as_mode());
 
     // Map row and column locations to byte slices (lazily).
-    let locator = Locator::new(contents);
+    let locator = Locator::new(source_kind.source_code());
 
     // Detect the current code style (lazily).
     let stylist = Stylist::from_tokens(&tokens, &locator);
@@ -374,7 +374,7 @@ pub fn lint_only(
         &directives,
         settings,
         noqa,
-        source_kind,
+        Some(source_kind),
         source_type,
     );
 
@@ -416,15 +416,14 @@ fn diagnostics_to_messages(
 /// Generate `Diagnostic`s from source code content, iteratively autofixing
 /// until stable.
 pub fn lint_fix<'a>(
-    contents: &'a str,
     path: &Path,
     package: Option<&Path>,
     noqa: flags::Noqa,
     settings: &Settings,
-    source_kind: &mut SourceKind,
+    source_kind: &'a SourceKind,
     source_type: PySourceType,
 ) -> Result<FixerResult<'a>> {
-    let mut transformed = Cow::Borrowed(contents);
+    let mut transformed = Cow::Borrowed(source_kind);
 
     // Track the number of fixed errors across iterations.
     let mut fixed = FxHashMap::default();
@@ -439,10 +438,10 @@ pub fn lint_fix<'a>(
     loop {
         // Tokenize once.
         let tokens: Vec<LexResult> =
-            ruff_python_parser::tokenize(&transformed, source_type.as_mode());
+            ruff_python_parser::tokenize(transformed.source_code(), source_type.as_mode());
 
         // Map row and column locations to byte slices (lazily).
-        let locator = Locator::new(&transformed);
+        let locator = Locator::new(transformed.source_code());
 
         // Detect the current code style (lazily).
         let stylist = Stylist::from_tokens(&tokens, &locator);
@@ -482,7 +481,7 @@ pub fn lint_fix<'a>(
             if parseable && result.error.is_some() {
                 report_autofix_syntax_error(
                     path,
-                    &transformed,
+                    transformed.source_code(),
                     &result.error.unwrap(),
                     fixed.keys().copied(),
                 );
@@ -503,12 +502,7 @@ pub fn lint_fix<'a>(
                     *fixed.entry(rule).or_default() += count;
                 }
 
-                if let SourceKind::Jupyter(notebook) = source_kind {
-                    notebook.update(&source_map, &fixed_contents);
-                }
-
-                // Store the fixed contents.
-                transformed = Cow::Owned(fixed_contents);
+                transformed = Cow::Owned(transformed.updated(fixed_contents, &source_map));
 
                 // Increment the iteration count.
                 iterations += 1;
@@ -517,7 +511,7 @@ pub fn lint_fix<'a>(
                 continue;
             }
 
-            report_failed_to_converge_error(path, &transformed, &result.data.0);
+            report_failed_to_converge_error(path, transformed.source_code(), &result.data.0);
         }
 
         return Ok(FixerResult {

crates/ruff/src/message/azure.rs

@@ -18,7 +18,7 @@ impl Emitter for AzureEmitter {
         context: &EmitterContext,
     ) -> anyhow::Result<()> {
         for message in messages {
-            let location = if context.is_jupyter_notebook(message.filename()) {
+            let location = if context.is_notebook(message.filename()) {
                 // We can't give a reasonable location for the structured formats,
                 // so we show one that's clearly a fallback
                 SourceLocation::default()

crates/ruff/src/message/github.rs

@@ -20,7 +20,7 @@ impl Emitter for GithubEmitter {
     ) -> anyhow::Result<()> {
         for message in messages {
             let source_location = message.compute_start_location();
-            let location = if context.is_jupyter_notebook(message.filename()) {
+            let location = if context.is_notebook(message.filename()) {
                 // We can't give a reasonable location for the structured formats,
                 // so we show one that's clearly a fallback
                 SourceLocation::default()

crates/ruff/src/message/gitlab.rs

@@ -63,7 +63,7 @@ impl Serialize for SerializedMessages<'_> {
             let start_location = message.compute_start_location();
             let end_location = message.compute_end_location();
 
-            let lines = if self.context.is_jupyter_notebook(message.filename()) {
+            let lines = if self.context.is_notebook(message.filename()) {
                 // We can't give a reasonable location for the structured formats,
                 // so we show one that's clearly a fallback
                 json!({

crates/ruff/src/message/grouped.rs

@@ -13,7 +13,6 @@ use crate::message::text::{MessageCodeFrame, RuleCodeAndBody};
 use crate::message::{
     group_messages_by_filename, Emitter, EmitterContext, Message, MessageWithLocation,
 };
-use crate::source_kind::SourceKind;
 
 #[derive(Default)]
 pub struct GroupedEmitter {
@@ -66,10 +65,7 @@ impl Emitter for GroupedEmitter {
                     writer,
                     "{}",
                     DisplayGroupedMessage {
-                        jupyter_index: context
-                            .source_kind(message.filename())
-                            .and_then(SourceKind::notebook)
-                            .map(Notebook::index),
+                        jupyter_index: context.notebook(message.filename()).map(Notebook::index),
                         message,
                         show_fix_status: self.show_fix_status,
                         show_source: self.show_source,

crates/ruff/src/message/junit.rs

@@ -45,7 +45,7 @@ impl Emitter for JunitEmitter {
                     } = message;
                     let mut status = TestCaseStatus::non_success(NonSuccessKind::Failure);
                     status.set_message(message.kind.body.clone());
-                    let location = if context.is_jupyter_notebook(message.filename()) {
+                    let location = if context.is_notebook(message.filename()) {
                         // We can't give a reasonable location for the structured formats,
                         // so we show one that's clearly a fallback
                         SourceLocation::default()

crates/ruff/src/message/mod.rs

@@ -3,10 +3,8 @@ use std::collections::BTreeMap;
 use std::io::Write;
 use std::ops::Deref;
 
-use ruff_text_size::{TextRange, TextSize};
 use rustc_hash::FxHashMap;
 
-use crate::source_kind::SourceKind;
 pub use azure::AzureEmitter;
 pub use github::GithubEmitter;
 pub use gitlab::GitlabEmitter;
@@ -17,8 +15,11 @@ pub use junit::JunitEmitter;
 pub use pylint::PylintEmitter;
 use ruff_diagnostics::{Diagnostic, DiagnosticKind, Fix};
 use ruff_source_file::{SourceFile, SourceLocation};
+use ruff_text_size::{TextRange, TextSize};
 pub use text::TextEmitter;
 
+use crate::jupyter::Notebook;
+
 mod azure;
 mod diff;
 mod github;
@@ -129,33 +130,31 @@ pub trait Emitter {
 
 /// Context passed to [`Emitter`].
 pub struct EmitterContext<'a> {
-    source_kind: &'a FxHashMap<String, SourceKind>,
+    notebooks: &'a FxHashMap<String, Notebook>,
 }
 
 impl<'a> EmitterContext<'a> {
-    pub fn new(source_kind: &'a FxHashMap<String, SourceKind>) -> Self {
-        Self { source_kind }
+    pub fn new(notebooks: &'a FxHashMap<String, Notebook>) -> Self {
+        Self { notebooks }
     }
 
     /// Tests if the file with `name` is a jupyter notebook.
-    pub fn is_jupyter_notebook(&self, name: &str) -> bool {
-        self.source_kind
-            .get(name)
-            .is_some_and(SourceKind::is_jupyter)
+    pub fn is_notebook(&self, name: &str) -> bool {
+        self.notebooks.contains_key(name)
     }
 
-    pub fn source_kind(&self, name: &str) -> Option<&SourceKind> {
-        self.source_kind.get(name)
+    pub fn notebook(&self, name: &str) -> Option<&Notebook> {
+        self.notebooks.get(name)
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use ruff_text_size::{TextRange, TextSize};
     use rustc_hash::FxHashMap;
 
     use ruff_diagnostics::{Diagnostic, DiagnosticKind, Edit, Fix};
     use ruff_source_file::SourceFileBuilder;
+    use ruff_text_size::{TextRange, TextSize};
 
     use crate::message::{Emitter, EmitterContext, Message};
 

crates/ruff/src/message/pylint.rs

@@ -19,7 +19,7 @@ impl Emitter for PylintEmitter {
         context: &EmitterContext,
     ) -> anyhow::Result<()> {
         for message in messages {
-            let row = if context.is_jupyter_notebook(message.filename()) {
+            let row = if context.is_notebook(message.filename()) {
                 // We can't give a reasonable location for the structured formats,
                 // so we show one that's clearly a fallback
                 OneIndexed::from_zero_indexed(0)

crates/ruff/src/message/text.rs

@@ -6,17 +6,16 @@ use annotate_snippets::display_list::{DisplayList, FormatOptions};
 use annotate_snippets::snippet::{Annotation, AnnotationType, Slice, Snippet, SourceAnnotation};
 use bitflags::bitflags;
 use colored::Colorize;
-use ruff_text_size::{TextRange, TextSize};
 
 use ruff_source_file::{OneIndexed, SourceLocation};
+use ruff_text_size::{TextRange, TextSize};
 
 use crate::fs::relativize_path;
 use crate::jupyter::{JupyterIndex, Notebook};
 use crate::line_width::{LineWidth, TabSize};
 use crate::message::diff::Diff;
 use crate::message::{Emitter, EmitterContext, Message};
 use crate::registry::AsRule;
-use crate::source_kind::SourceKind;
 
 bitflags! {
     #[derive(Default)]
@@ -72,10 +71,7 @@ impl Emitter for TextEmitter {
             )?;
 
             let start_location = message.compute_start_location();
-            let jupyter_index = context
-                .source_kind(message.filename())
-                .and_then(SourceKind::notebook)
-                .map(Notebook::index);
+            let jupyter_index = context.notebook(message.filename()).map(Notebook::index);
 
             // Check if we're working on a jupyter notebook and translate positions with cell accordingly
             let diagnostic_location = if let Some(jupyter_index) = jupyter_index {