Store source code on Message

Author: MichaReiser

Cargo.lock

@@ -21,6 +21,7 @@ ruff_python_ast = { path = "../ruff_python_ast" }
 ruff_python_semantic = { path = "../ruff_python_semantic" }
 ruff_python_stdlib = { path = "../ruff_python_stdlib" }
 ruff_rustpython = { path = "../ruff_rustpython" }
+ruff_text_size = { path = "../ruff_text_size" }
 
 annotate-snippets = { version = "0.9.1", features = ["color"] }
 anyhow = { workspace = true }

crates/ruff/Cargo.toml

@@ -103,7 +103,7 @@ fn is_lone_child(child: &Stmt, parent: &Stmt, deleted: &[&Stmt]) -> Result<bool>
 /// Return the location of a trailing semicolon following a `Stmt`, if it's part
 /// of a multi-statement line.
 fn trailing_semicolon(stmt: &Stmt, locator: &Locator) -> Option<Location> {
-    let contents = locator.skip(stmt.end_location.unwrap());
+    let contents = locator.after(stmt.end_location.unwrap());
     for (row, line) in NewlineWithTrailingNewline::from(contents).enumerate() {
         let trimmed = line.trim();
         if trimmed.starts_with(';') {
@@ -126,7 +126,7 @@ fn trailing_semicolon(stmt: &Stmt, locator: &Locator) -> Option<Location> {
 /// Find the next valid break for a `Stmt` after a semicolon.
 fn next_stmt_break(semicolon: Location, locator: &Locator) -> Location {
     let start_location = Location::new(semicolon.row(), semicolon.column() + 1);
-    let contents = locator.skip(start_location);
+    let contents = locator.after(start_location);
     for (row, line) in NewlineWithTrailingNewline::from(contents).enumerate() {
         let trimmed = line.trim();
         // Skip past any continuations.
@@ -158,7 +158,7 @@ fn next_stmt_break(semicolon: Location, locator: &Locator) -> Location {
 
 /// Return `true` if a `Stmt` occurs at the end of a file.
 fn is_end_of_file(stmt: &Stmt, locator: &Locator) -> bool {
-    let contents = locator.skip(stmt.end_location.unwrap());
+    let contents = locator.after(stmt.end_location.unwrap());
     contents.is_empty()
 }
 
@@ -361,7 +361,7 @@ pub fn remove_argument(
     remove_parentheses: bool,
 ) -> Result<Edit> {
     // TODO(sbrugman): Preserve trailing comments.
-    let contents = locator.skip(call_at);
+    let contents = locator.after(call_at);
 
     let mut fix_start = None;
     let mut fix_end = None;

crates/ruff/src/autofix/actions.rs

@@ -80,7 +80,7 @@ fn apply_fixes<'a>(
     }
 
     // Add the remaining content.
-    let slice = locator.skip(last_pos.unwrap_or_default());
+    let slice = locator.after(last_pos.unwrap_or_default());
     output.push_str(slice);
 
     (output, fixed)

crates/ruff/src/autofix/mod.rs

@@ -174,7 +174,7 @@ fn match_docstring_end(body: &[Stmt]) -> Option<Location> {
 /// along with a trailing newline suffix.
 fn end_of_statement_insertion(stmt: &Stmt, locator: &Locator, stylist: &Stylist) -> Insertion {
     let location = stmt.end_location.unwrap();
-    let mut tokens = lexer::lex_located(locator.skip(location), Mode::Module, location).flatten();
+    let mut tokens = lexer::lex_located(locator.after(location), Mode::Module, location).flatten();
     if let Some((.., Tok::Semi, end)) = tokens.next() {
         // If the first token after the docstring is a semicolon, insert after the semicolon as an
         // inline statement;
@@ -207,7 +207,7 @@ fn top_of_file_insertion(body: &[Stmt], locator: &Locator, stylist: &Stylist) ->
     let mut location = if let Some(location) = match_docstring_end(body) {
         // If the first token after the docstring is a semicolon, insert after the semicolon as an
         // inline statement;
-        let first_token = lexer::lex_located(locator.skip(location), Mode::Module, location)
+        let first_token = lexer::lex_located(locator.after(location), Mode::Module, location)
             .flatten()
             .next();
         if let Some((.., Tok::Semi, end)) = first_token {
@@ -222,7 +222,7 @@ fn top_of_file_insertion(body: &[Stmt], locator: &Locator, stylist: &Stylist) ->
 
     // Skip over any comments and empty lines.
     for (.., tok, end) in
-        lexer::lex_located(locator.skip(location), Mode::Module, location).flatten()
+        lexer::lex_located(locator.after(location), Mode::Module, location).flatten()
     {
         if matches!(tok, Tok::Comment(..) | Tok::Newline) {
             location = Location::new(end.row() + 1, 0);

crates/ruff/src/importer.rs

@@ -22,7 +22,7 @@ use crate::checkers::physical_lines::check_physical_lines;
 use crate::checkers::tokens::check_tokens;
 use crate::directives::Directives;
 use crate::doc_lines::{doc_lines_from_ast, doc_lines_from_tokens};
-use crate::message::{Message, Source};
+use crate::message::Message;
 use crate::noqa::add_noqa;
 use crate::registry::{AsRule, Rule};
 use crate::rules::pycodestyle;
@@ -355,19 +355,26 @@ pub fn lint_only(
 
     // Convert from diagnostics to messages.
     let path_lossy = path.to_string_lossy();
+
     result.map(|(messages, imports)| {
+        let source_code = if settings.show_source && !messages.is_empty() {
+            Some(locator.to_source_code_buf())
+        } else {
+            None
+        };
+
         (
             messages
                 .into_iter()
                 .map(|diagnostic| {
-                    let source = if settings.show_source {
-                        Some(Source::from_diagnostic(&diagnostic, &locator))
-                    } else {
-                        None
-                    };
                     let lineno = diagnostic.location.row();
                     let noqa_row = *directives.noqa_line_for.get(&lineno).unwrap_or(&lineno);
-                    Message::from_diagnostic(diagnostic, path_lossy.to_string(), source, noqa_row)
+                    Message::from_diagnostic(
+                        diagnostic,
+                        path_lossy.to_string(),
+                        source_code.clone(),
+                        noqa_row,
+                    )
                 })
                 .collect(),
             imports,
@@ -500,22 +507,23 @@ This indicates a bug in `{}`. If you could open an issue at:
         let path_lossy = path.to_string_lossy();
         return Ok(FixerResult {
             result: result.map(|(messages, imports)| {
+                let source_code = if settings.show_source && !messages.is_empty() {
+                    Some(locator.to_source_code_buf())
+                } else {
+                    None
+                };
+
                 (
                     messages
                         .into_iter()
                         .map(|diagnostic| {
-                            let source = if settings.show_source {
-                                Some(Source::from_diagnostic(&diagnostic, &locator))
-                            } else {
-                                None
-                            };
                             let lineno = diagnostic.location.row();
                             let noqa_row =
                                 *directives.noqa_line_for.get(&lineno).unwrap_or(&lineno);
                             Message::from_diagnostic(
                                 diagnostic,
                                 path_lossy.to_string(),
-                                source,
+                                source_code.clone(),
                                 noqa_row,
                             )
                         })

crates/ruff/src/linter.rs

@@ -1,12 +1,9 @@
 use crate::fs::relativize_path;
 use crate::jupyter::JupyterIndex;
-use crate::message::text::RuleCodeAndBody;
+use crate::message::text::{MessageCodeFrame, RuleCodeAndBody};
 use crate::message::{group_messages_by_filename, Emitter, EmitterContext, Message};
-use crate::registry::AsRule;
-use annotate_snippets::display_list::{DisplayList, FormatOptions};
-use annotate_snippets::snippet::{Annotation, AnnotationType, Slice, Snippet, SourceAnnotation};
 use colored::Colorize;
-use std::fmt::{Display, Formatter, Write as fmtWrite};
+use std::fmt::{Display, Formatter};
 use std::io::Write;
 
 #[derive(Default)]
@@ -116,47 +113,10 @@ impl Display for DisplayGroupedMessage<'_> {
             },
         )?;
 
-        if let Some(source) = &message.source {
-            let suggestion = message.kind.suggestion.clone();
-            let footer = if suggestion.is_some() {
-                vec![Annotation {
-                    id: None,
-                    label: suggestion.as_deref(),
-                    annotation_type: AnnotationType::Help,
-                }]
-            } else {
-                vec![]
-            };
-            let label = message.kind.rule().noqa_code().to_string();
-            let snippet = Snippet {
-                title: None,
-                footer,
-                slices: vec![Slice {
-                    source: &source.contents,
-                    line_start: message.location.row(),
-                    annotations: vec![SourceAnnotation {
-                        label: &label,
-                        annotation_type: AnnotationType::Error,
-                        range: source.range,
-                    }],
-                    // The origin (file name, line number, and column number) is already encoded
-                    // in the `label`.
-                    origin: None,
-                    fold: false,
-                }],
-                opt: FormatOptions {
-                    #[cfg(test)]
-                    color: false,
-                    #[cfg(not(test))]
-                    color: colored::control::SHOULD_COLORIZE.should_colorize(),
-                    ..FormatOptions::default()
-                },
-            };
-            // Skip the first line, since we format the `label` ourselves.
-            let message = DisplayList::from(snippet);
+        {
+            use std::fmt::Write;
             let mut padded = PadAdapter::new(f);
-
-            writeln!(&mut padded, "{message}")?;
+            write!(padded, "{}", MessageCodeFrame { message })?;
         }
 
         writeln!(f)?;

crates/ruff/src/message/grouped.rs

@@ -20,30 +20,28 @@ pub use json::JsonEmitter;
 pub use junit::JunitEmitter;
 pub use pylint::PylintEmitter;
 pub use rustpython_parser::ast::Location;
-use serde::{Deserialize, Serialize};
 pub use text::TextEmitter;
 
 use crate::jupyter::JupyterIndex;
 use ruff_diagnostics::{Diagnostic, DiagnosticKind, Fix};
-use ruff_python_ast::source_code::Locator;
-use ruff_python_ast::types::Range;
+use ruff_python_ast::source_code::SourceCodeBuf;
 
-#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, PartialEq, Eq)]
 pub struct Message {
     pub kind: DiagnosticKind,
     pub location: Location,
     pub end_location: Location,
     pub fix: Fix,
     pub filename: String,
-    pub source: Option<Source>,
+    pub source: Option<SourceCodeBuf>,
     pub noqa_row: usize,
 }
 
 impl Message {
     pub fn from_diagnostic(
         diagnostic: Diagnostic,
         filename: String,
-        source: Option<Source>,
+        source: Option<SourceCodeBuf>,
         noqa_row: usize,
     ) -> Self {
         Self {
@@ -77,38 +75,6 @@ impl PartialOrd for Message {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone)]
-pub struct Source {
-    pub contents: String,
-    pub range: (usize, usize),
-}
-
-impl Source {
-    pub fn from_diagnostic(diagnostic: &Diagnostic, locator: &Locator) -> Self {
-        let location = Location::new(diagnostic.location.row(), 0);
-        // Diagnostics can already extend one-past-the-end. If they do, though, then
-        // they'll end at the start of a line. We need to avoid extending by yet another
-        // line past-the-end.
-        let end_location = if diagnostic.end_location.column() == 0 {
-            diagnostic.end_location
-        } else {
-            Location::new(diagnostic.end_location.row() + 1, 0)
-        };
-        let source = locator.slice(Range::new(location, end_location));
-        let num_chars_in_range = locator
-            .slice(Range::new(diagnostic.location, diagnostic.end_location))
-            .chars()
-            .count();
-        Source {
-            contents: source.to_string(),
-            range: (
-                diagnostic.location.column(),
-                diagnostic.location.column() + num_chars_in_range,
-            ),
-        }
-    }
-}
-
 fn group_messages_by_filename(messages: &[Message]) -> BTreeMap<&str, Vec<&Message>> {
     let mut grouped_messages = BTreeMap::default();
     for message in messages {
@@ -156,15 +122,15 @@ impl<'a> EmitterContext<'a> {
 
 #[cfg(test)]
 mod tests {
-    use crate::message::{Emitter, EmitterContext, Location, Message, Source};
+    use crate::message::{Emitter, EmitterContext, Location, Message};
     use crate::rules::pyflakes::rules::{UndefinedName, UnusedImport, UnusedVariable};
     use ruff_diagnostics::{Diagnostic, Edit, Fix};
-    use ruff_python_ast::source_code::Locator;
+    use ruff_python_ast::source_code::SourceCodeBuf;
     use ruff_python_ast::types::Range;
     use rustc_hash::FxHashMap;
 
     pub(super) fn create_messages() -> Vec<Message> {
-        let file_1 = r#"import os
+        let fib = r#"import os
 
 def fibonacci(n):
     """Compute the nth number in the Fibonacci sequence."""
@@ -177,8 +143,6 @@ def fibonacci(n):
         return fibonacci(n - 1) + fibonacci(n - 2)
         "#;
 
-        let file_1_locator = Locator::new(file_1);
-
         let unused_import = Diagnostic::new(
             UnusedImport {
                 name: "os".to_string(),
@@ -188,7 +152,7 @@ def fibonacci(n):
             Range::new(Location::new(1, 7), Location::new(1, 9)),
         );
 
-        let unused_source = Source::from_diagnostic(&unused_import, &file_1_locator);
+        let fib_source = SourceCodeBuf::from_content(fib);
 
         let unused_variable = Diagnostic::new(
             UnusedVariable {
@@ -201,8 +165,6 @@ def fibonacci(n):
             Location::new(5, 9),
         )]));
 
-        let unused_variable_source = Source::from_diagnostic(&unused_variable, &file_1_locator);
-
         let file_2 = r#"if a == 1: pass"#;
 
         let undefined_name = Diagnostic::new(
@@ -212,20 +174,20 @@ def fibonacci(n):
             Range::new(Location::new(1, 3), Location::new(1, 4)),
         );
 
-        let undefined_source = Source::from_diagnostic(&undefined_name, &Locator::new(file_2));
+        let file_2_source = SourceCodeBuf::from_content(file_2);
 
         vec![
-            Message::from_diagnostic(unused_import, "fib.py".to_string(), Some(unused_source), 1),
             Message::from_diagnostic(
-                unused_variable,
+                unused_import,
                 "fib.py".to_string(),
-                Some(unused_variable_source),
+                Some(fib_source.clone()),
                 1,
             ),
+            Message::from_diagnostic(unused_variable, "fib.py".to_string(), Some(fib_source), 1),
             Message::from_diagnostic(
                 undefined_name,
                 "undef.py".to_string(),
-                Some(undefined_source),
+                Some(file_2_source),
                 1,
             ),
         ]