Render code frame with context (#3901)

crates/ruff/src/message/mod.rs

@@ -132,6 +132,7 @@ mod tests {
     pub(super) fn create_messages() -> Vec<Message> {
         let fib = r#"import os
 
+
 def fibonacci(n):
     """Compute the nth number in the Fibonacci sequence."""
     x = 1
@@ -141,7 +142,7 @@ def fibonacci(n):
         return 1
     else:
         return fibonacci(n - 1) + fibonacci(n - 2)
-        "#;
+"#;
 
         let unused_import = Diagnostic::new(
             UnusedImport {
@@ -158,11 +159,11 @@ def fibonacci(n):
             UnusedVariable {
                 name: "x".to_string(),
             },
-            Range::new(Location::new(5, 4), Location::new(5, 5)),
+            Range::new(Location::new(6, 4), Location::new(6, 5)),
         )
         .with_fix(Fix::new(vec![Edit::deletion(
-            Location::new(5, 4),
+            Location::new(6, 4),
-            Location::new(5, 9),
+            Location::new(6, 9),
         )]));
 
         let file_2 = r#"if a == 1: pass"#;

crates/ruff/src/message/snapshots/ruff__message__azure__tests__output.snap

@@ -3,6 +3,6 @@ source: crates/ruff/src/message/azure.rs
 expression: content
 ---
 ##vso[task.logissue type=error;sourcepath=fib.py;linenumber=1;columnnumber=8;code=F401;]`os` imported but unused
-##vso[task.logissue type=error;sourcepath=fib.py;linenumber=5;columnnumber=5;code=F841;]Local variable `x` is assigned to but never used
+##vso[task.logissue type=error;sourcepath=fib.py;linenumber=6;columnnumber=5;code=F841;]Local variable `x` is assigned to but never used
 ##vso[task.logissue type=error;sourcepath=undef.py;linenumber=1;columnnumber=4;code=F821;]Undefined name `a`
 

crates/ruff/src/message/snapshots/ruff__message__github__tests__output.snap

@@ -3,6 +3,6 @@ source: crates/ruff/src/message/github.rs
 expression: content
 ---
 ::error title=Ruff (F401),file=fib.py,line=1,col=8,endLine=1,endColumn=10::fib.py:1:8: F401 `os` imported but unused
-::error title=Ruff (F841),file=fib.py,line=5,col=5,endLine=5,endColumn=6::fib.py:5:5: F841 Local variable `x` is assigned to but never used
+::error title=Ruff (F841),file=fib.py,line=6,col=5,endLine=6,endColumn=6::fib.py:6:5: F841 Local variable `x` is assigned to but never used
 ::error title=Ruff (F821),file=undef.py,line=1,col=4,endLine=1,endColumn=5::undef.py:1:4: F821 Undefined name `a`
 

crates/ruff/src/message/snapshots/ruff__message__gitlab__tests__output.snap

@@ -1,6 +1,6 @@
 ---
 source: crates/ruff/src/message/gitlab.rs
-expression: output
+expression: redact_fingerprint(&content)
 ---
 [
   {
@@ -22,8 +22,8 @@ expression: output
     "location": {
       "path": "fib.py",
       "lines": {
-        "begin": 5,
+        "begin": 6,
-        "end": 5
+        "end": 6
       }
     }
   },

crates/ruff/src/message/snapshots/ruff__message__grouped__tests__default.snap

@@ -10,10 +10,14 @@ fib.py:
     |
     = help: Remove unused import: `os`
 
-  5:5  F841 Local variable `x` is assigned to but never used
+  6:5 F841 Local variable `x` is assigned to but never used
      |
-  5 |     x = 1
+   6 | def fibonacci(n):
+   7 |     """Compute the nth number in the Fibonacci sequence."""
+   8 |     x = 1
      |     ^ F841
+   9 |     if n == 0:
+  10 |         return 0
      |
      = help: Remove assignment to unused variable `x`
 

crates/ruff/src/message/snapshots/ruff__message__grouped__tests__fix_status.snap

@@ -10,10 +10,14 @@ fib.py:
     |
     = help: Remove unused import: `os`
 
-  5:5  F841 [*] Local variable `x` is assigned to but never used
+  6:5 F841 [*] Local variable `x` is assigned to but never used
      |
-  5 |     x = 1
+   6 | def fibonacci(n):
+   7 |     """Compute the nth number in the Fibonacci sequence."""
+   8 |     x = 1
      |     ^ F841
+   9 |     if n == 0:
+  10 |         return 0
      |
      = help: Remove assignment to unused variable `x`
 

crates/ruff/src/message/snapshots/ruff__message__json__tests__output.snap

@@ -27,22 +27,22 @@ expression: content
         {
           "content": "",
           "location": {
-            "row": 5,
+            "row": 6,
             "column": 4
           },
           "end_location": {
-            "row": 5,
+            "row": 6,
             "column": 9
           }
         }
       ]
     },
     "location": {
-      "row": 5,
+      "row": 6,
       "column": 5
     },
     "end_location": {
-      "row": 5,
+      "row": 6,
       "column": 6
     },
     "filename": "fib.py",

crates/ruff/src/message/snapshots/ruff__message__junit__tests__output.snap

@@ -8,8 +8,8 @@ expression: content
         <testcase name="org.ruff.F401" classname="fib" line="1" column="8">
             <failure message="`os` imported but unused">line 1, col 8, `os` imported but unused</failure>
         </testcase>
-        <testcase name="org.ruff.F841" classname="fib" line="5" column="5">
+        <testcase name="org.ruff.F841" classname="fib" line="6" column="5">
-            <failure message="Local variable `x` is assigned to but never used">line 5, col 5, Local variable `x` is assigned to but never used</failure>
+            <failure message="Local variable `x` is assigned to but never used">line 6, col 5, Local variable `x` is assigned to but never used</failure>
         </testcase>
     </testsuite>
     <testsuite name="undef.py" tests="1" disabled="0" errors="0" failures="1" package="org.ruff">

crates/ruff/src/message/snapshots/ruff__message__pylint__tests__output.snap

@@ -3,6 +3,6 @@ source: crates/ruff/src/message/pylint.rs
 expression: content
 ---
 fib.py:1: [F401] `os` imported but unused
-fib.py:5: [F841] Local variable `x` is assigned to but never used
+fib.py:6: [F841] Local variable `x` is assigned to but never used
 undef.py:1: [F821] Undefined name `a`
 

crates/ruff/src/message/snapshots/ruff__message__text__tests__default.snap

@@ -9,10 +9,14 @@ fib.py:1:8: F401 `os` imported but unused
   |
   = help: Remove unused import: `os`
 
-fib.py:5:5: F841 Local variable `x` is assigned to but never used
+fib.py:6:5: F841 Local variable `x` is assigned to but never used
    |
-5 |     x = 1
+ 6 | def fibonacci(n):
+ 7 |     """Compute the nth number in the Fibonacci sequence."""
+ 8 |     x = 1
    |     ^ F841
+ 9 |     if n == 0:
+10 |         return 0
    |
    = help: Remove assignment to unused variable `x`
 

crates/ruff/src/message/snapshots/ruff__message__text__tests__fix_status.snap

@@ -9,10 +9,14 @@ fib.py:1:8: F401 [*] `os` imported but unused
   |
   = help: Remove unused import: `os`
 
-fib.py:5:5: F841 [*] Local variable `x` is assigned to but never used
+fib.py:6:5: F841 [*] Local variable `x` is assigned to but never used
    |
-5 |     x = 1
+ 6 | def fibonacci(n):
+ 7 |     """Compute the nth number in the Fibonacci sequence."""
+ 8 |     x = 1
    |     ^ F841
+ 9 |     if n == 0:
+10 |         return 0
    |
    = help: Remove assignment to unused variable `x`
 

crates/ruff/src/message/text.rs

@@ -1,13 +1,13 @@
 use crate::fs::relativize_path;
-use crate::message::{Emitter, EmitterContext, Location, Message};
+use crate::message::{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 ruff_diagnostics::DiagnosticKind;
 use ruff_python_ast::source_code::OneIndexed;
-use ruff_python_ast::types::Range;
 use ruff_text_size::TextRange;
+use std::cmp;
 use std::fmt::{Display, Formatter};
 use std::io::Write;
 
@@ -139,29 +139,52 @@ impl Display for MessageCodeFrame<'_> {
                 Vec::new()
             };
 
-            let source_code_start =
+            let mut start_index =
-                source_code.line_start(OneIndexed::new(location.row()).unwrap());
+                OneIndexed::new(cmp::max(1, location.row().saturating_sub(2))).unwrap();
+            let content_start_index = OneIndexed::new(location.row()).unwrap();
 
-            let source_code_end = source_code.line_start(
+            // Trim leading empty lines.
-                OneIndexed::new(
+            while start_index < content_start_index {
-                    end_location
+                if !source_code.line_text(start_index).trim().is_empty() {
-                        .row()
+                    break;
-                        .saturating_add(1)
+                }
-                        .min(source_code.line_count() + 1),
+                start_index = start_index.saturating_add(1);
-                )
+            }
-                .unwrap(),
-            );
 
-            let source_text =
+            let mut end_index = OneIndexed::new(cmp::min(
-                &source_code.text()[TextRange::new(source_code_start, source_code_end)];
+                end_location.row().saturating_add(2),
+                source_code.line_count() + 1,
+            ))
+            .unwrap();
 
-            let content_range = source_code.text_range(Range::new(
+            let content_end_index = OneIndexed::new(end_location.row()).unwrap();
-                // Subtract 1 because message column indices are 1 based but the index columns are 1 based.
-                Location::new(location.row(), location.column().saturating_sub(1)),
-                Location::new(end_location.row(), end_location.column().saturating_sub(1)),
-            ));
 
-            let annotation_length = &source_text[content_range - source_code_start]
+            // Trim trailing empty lines
+            while end_index > content_end_index {
+                if !source_code.line_text(end_index).trim().is_empty() {
+                    break;
+                }
+
+                end_index = end_index.saturating_sub(1);
+            }
+
+            let start_offset = source_code.line_start(start_index);
+            let end_offset = source_code.line_end(end_index);
+
+            let source_text = &source_code.text()[TextRange::new(start_offset, end_offset)];
+
+            let annotation_start_offset =
+                // Message columns are one indexed
+                source_code.offset(location.with_col_offset(-1)) - start_offset;
+            let annotation_end_offset =
+                source_code.offset(end_location.with_col_offset(-1)) - start_offset;
+
+            let start_char = source_text[TextRange::up_to(annotation_start_offset)]
+                .chars()
+                .count();
+
+            let char_length = source_text
+                [TextRange::new(annotation_start_offset, annotation_end_offset)]
             .chars()
             .count();
 
@@ -175,10 +198,7 @@ impl Display for MessageCodeFrame<'_> {
                     annotations: vec![SourceAnnotation {
                         label: &label,
                         annotation_type: AnnotationType::Error,
-                        range: (
+                        range: (start_char, start_char + char_length),
-                            location.column() - 1,
-                            location.column() + annotation_length - 1,
-                        ),
                     }],
                     // The origin (file name, line number, and column number) is already encoded
                     // in the `label`.

crates/ruff_python_ast/src/source_code/line_index.rs

@@ -100,6 +100,20 @@ impl LineIndex {
         }
     }
 
+    /// Returns the [byte offset](TextSize) of the `line`'s end.
+    /// The offset is the end of the line, up to and including the newline character ending the line (if any).
+    pub(crate) fn line_end(&self, line: OneIndexed, contents: &str) -> TextSize {
+        let row_index = line.to_zero_indexed();
+        let starts = self.line_starts();
+
+        // If start-of-line position after last line
+        if row_index.saturating_add(1) >= starts.len() {
+            contents.text_len()
+        } else {
+            starts[row_index + 1]
+        }
+    }
+
     /// Returns the [`TextRange`] of the `line` with the given index.
     /// The start points to the first character's [byte offset](TextSize), the end up to, and including
     /// the newline character ending the line (if any).

crates/ruff_python_ast/src/source_code/mod.rs

@@ -78,6 +78,10 @@ impl<'src, 'index> SourceCode<'src, 'index> {
         self.index.line_start(line, self.text)
     }
 
+    pub fn line_end(&self, line: OneIndexed) -> TextSize {
+        self.index.line_end(line, self.text)
+    }
+
     pub fn line_range(&self, line: OneIndexed) -> TextRange {
         self.index.line_range(line, self.text)
     }
@@ -198,6 +202,11 @@ impl SourceCodeBuf {
         self.as_source_code().offset(location)
     }
 
+    #[inline]
+    pub fn line_end(&self, line: OneIndexed) -> TextSize {
+        self.as_source_code().line_end(line)
+    }
+
     #[inline]
     pub fn line_start(&self, line: OneIndexed) -> TextSize {
         self.as_source_code().line_start(line)