[red-knot] ruff_db: make diagnostic rendering prettier (#15856)

This change does a simple swap of the existing renderer for one that
uses our vendored copy of `annotate-snippets`. We don't change anything
about the diagnostic data model, but this alone already makes
diagnostics look a lot nicer!

Cargo.lock

@@ -2717,6 +2717,7 @@ name = "ruff_db"
 version = "0.0.0"
 dependencies = [
  "camino",
+ "colored 3.0.0",
  "countme",
  "dashmap 6.1.0",
  "dunce",
@@ -2726,6 +2727,7 @@ dependencies = [
  "insta",
  "matchit",
  "path-slash",
+ "ruff_annotate_snippets",
  "ruff_cache",
  "ruff_notebook",
  "ruff_python_ast",

crates/red_knot/tests/cli.rs

@@ -28,14 +28,21 @@ fn config_override() -> anyhow::Result<()> {
         ),
     ])?;
 
-    assert_cmd_snapshot!(case.command(), @r"
+    assert_cmd_snapshot!(case.command(), @r###"
-        success: false
+    success: false
-        exit_code: 1
+    exit_code: 1
-        ----- stdout -----
+    ----- stdout -----
-        error[lint:unresolved-attribute] <temp_dir>/test.py:5:7 Type `<module 'sys'>` has no attribute `last_exc`
+    error: lint:unresolved-attribute
+     --> <temp_dir>/test.py:5:7
+      |
+    4 | # Access `sys.last_exc` that was only added in Python 3.12
+    5 | print(sys.last_exc)
+      |       ^^^^^^^^^^^^ Type `<module 'sys'>` has no attribute `last_exc`
+      |
 
-        ----- stderr -----
+
-    ");
+    ----- stderr -----
+    "###);
 
     assert_cmd_snapshot!(case.command().arg("--python-version").arg("3.12"), @r"
         success: true
@@ -91,14 +98,22 @@ fn cli_arguments_are_relative_to_the_current_directory() -> anyhow::Result<()> {
     ])?;
 
     // Make sure that the CLI fails when the `libs` directory is not in the search path.
-    assert_cmd_snapshot!(case.command().current_dir(case.project_dir().join("child")), @r#"
+    assert_cmd_snapshot!(case.command().current_dir(case.project_dir().join("child")), @r###"
-        success: false
+    success: false
-        exit_code: 1
+    exit_code: 1
-        ----- stdout -----
+    ----- stdout -----
-        error[lint:unresolved-import] <temp_dir>/child/test.py:2:1 Cannot resolve import `utils`
+    error: lint:unresolved-import
+     --> <temp_dir>/child/test.py:2:1
+      |
+    2 | from utils import add
+      | ^^^^^^^^^^^^^^^^^^^^^ Cannot resolve import `utils`
+    3 |
+    4 | stat = add(10, 15)
+      |
 
-        ----- stderr -----
+
-    "#);
+    ----- stderr -----
+    "###);
 
     assert_cmd_snapshot!(case.command().current_dir(case.project_dir().join("child")).arg("--extra-search-path").arg("../libs"), @r"
         success: true
@@ -180,15 +195,31 @@ fn configuration_rule_severity() -> anyhow::Result<()> {
 
     // Assert that there's a possibly unresolved reference diagnostic
     // and that division-by-zero has a severity of error by default.
-    assert_cmd_snapshot!(case.command(), @r"
+    assert_cmd_snapshot!(case.command(), @r###"
-        success: false
+    success: false
-        exit_code: 1
+    exit_code: 1
-        ----- stdout -----
+    ----- stdout -----
-        error[lint:division-by-zero] <temp_dir>/test.py:2:5 Cannot divide object of type `Literal[4]` by zero
+    error: lint:division-by-zero
-        warning[lint:possibly-unresolved-reference] <temp_dir>/test.py:7:7 Name `x` used when possibly not defined
+     --> <temp_dir>/test.py:2:5
+      |
+    2 | y = 4 / 0
+      |     ^^^^^ Cannot divide object of type `Literal[4]` by zero
+    3 |
+    4 | for a in range(0, y):
+      |
 
-        ----- stderr -----
+    warning: lint:possibly-unresolved-reference
-    ");
+     --> <temp_dir>/test.py:7:7
+      |
+    5 |     x = a
+    6 |
+    7 | print(x)  # possibly-unresolved-reference
+      |       - Name `x` used when possibly not defined
+      |
+
+
+    ----- stderr -----
+    "###);
 
     case.write_file(
         "pyproject.toml",
@@ -199,14 +230,22 @@ fn configuration_rule_severity() -> anyhow::Result<()> {
     "#,
     )?;
 
-    assert_cmd_snapshot!(case.command(), @r"
+    assert_cmd_snapshot!(case.command(), @r###"
-        success: false
+    success: false
-        exit_code: 1
+    exit_code: 1
-        ----- stdout -----
+    ----- stdout -----
-        warning[lint:division-by-zero] <temp_dir>/test.py:2:5 Cannot divide object of type `Literal[4]` by zero
+    warning: lint:division-by-zero
+     --> <temp_dir>/test.py:2:5
+      |
+    2 | y = 4 / 0
+      |     ----- Cannot divide object of type `Literal[4]` by zero
+    3 |
+    4 | for a in range(0, y):
+      |
 
-        ----- stderr -----
+
-    ");
+    ----- stderr -----
+    "###);
 
     Ok(())
 }
@@ -230,16 +269,42 @@ fn cli_rule_severity() -> anyhow::Result<()> {
 
     // Assert that there's a possibly unresolved reference diagnostic
     // and that division-by-zero has a severity of error by default.
-    assert_cmd_snapshot!(case.command(), @r"
+    assert_cmd_snapshot!(case.command(), @r###"
-        success: false
+    success: false
-        exit_code: 1
+    exit_code: 1
-        ----- stdout -----
+    ----- stdout -----
-        error[lint:unresolved-import] <temp_dir>/test.py:2:8 Cannot resolve import `does_not_exit`
+    error: lint:unresolved-import
-        error[lint:division-by-zero] <temp_dir>/test.py:4:5 Cannot divide object of type `Literal[4]` by zero
+     --> <temp_dir>/test.py:2:8
-        warning[lint:possibly-unresolved-reference] <temp_dir>/test.py:9:7 Name `x` used when possibly not defined
+      |
+    2 | import does_not_exit
+      |        ^^^^^^^^^^^^^ Cannot resolve import `does_not_exit`
+    3 |
+    4 | y = 4 / 0
+      |
 
-        ----- stderr -----
+    error: lint:division-by-zero
-    ");
+     --> <temp_dir>/test.py:4:5
+      |
+    2 | import does_not_exit
+    3 |
+    4 | y = 4 / 0
+      |     ^^^^^ Cannot divide object of type `Literal[4]` by zero
+    5 |
+    6 | for a in range(0, y):
+      |
+
+    warning: lint:possibly-unresolved-reference
+     --> <temp_dir>/test.py:9:7
+      |
+    7 |     x = a
+    8 |
+    9 | print(x)  # possibly-unresolved-reference
+      |       - Name `x` used when possibly not defined
+      |
+
+
+    ----- stderr -----
+    "###);
 
     assert_cmd_snapshot!(
         case
@@ -250,15 +315,33 @@ fn cli_rule_severity() -> anyhow::Result<()> {
             .arg("division-by-zero")
             .arg("--warn")
             .arg("unresolved-import"),
-        @r"
+        @r###"
     success: false
     exit_code: 1
     ----- stdout -----
-    warning[lint:unresolved-import] <temp_dir>/test.py:2:8 Cannot resolve import `does_not_exit`
+    warning: lint:unresolved-import
-    warning[lint:division-by-zero] <temp_dir>/test.py:4:5 Cannot divide object of type `Literal[4]` by zero
+     --> <temp_dir>/test.py:2:8
+      |
+    2 | import does_not_exit
+      |        ------------- Cannot resolve import `does_not_exit`
+    3 |
+    4 | y = 4 / 0
+      |
+
+    warning: lint:division-by-zero
+     --> <temp_dir>/test.py:4:5
+      |
+    2 | import does_not_exit
+    3 |
+    4 | y = 4 / 0
+      |     ----- Cannot divide object of type `Literal[4]` by zero
+    5 |
+    6 | for a in range(0, y):
+      |
+
 
     ----- stderr -----
-    "
+    "###
     );
 
     Ok(())
@@ -282,15 +365,31 @@ fn cli_rule_severity_precedence() -> anyhow::Result<()> {
 
     // Assert that there's a possibly unresolved reference diagnostic
     // and that division-by-zero has a severity of error by default.
-    assert_cmd_snapshot!(case.command(), @r"
+    assert_cmd_snapshot!(case.command(), @r###"
-        success: false
+    success: false
-        exit_code: 1
+    exit_code: 1
-        ----- stdout -----
+    ----- stdout -----
-        error[lint:division-by-zero] <temp_dir>/test.py:2:5 Cannot divide object of type `Literal[4]` by zero
+    error: lint:division-by-zero
-        warning[lint:possibly-unresolved-reference] <temp_dir>/test.py:7:7 Name `x` used when possibly not defined
+     --> <temp_dir>/test.py:2:5
+      |
+    2 | y = 4 / 0
+      |     ^^^^^ Cannot divide object of type `Literal[4]` by zero
+    3 |
+    4 | for a in range(0, y):
+      |
 
-        ----- stderr -----
+    warning: lint:possibly-unresolved-reference
-    ");
+     --> <temp_dir>/test.py:7:7
+      |
+    5 |     x = a
+    6 |
+    7 | print(x)  # possibly-unresolved-reference
+      |       - Name `x` used when possibly not defined
+      |
+
+
+    ----- stderr -----
+    "###);
 
     assert_cmd_snapshot!(
         case
@@ -302,14 +401,22 @@ fn cli_rule_severity_precedence() -> anyhow::Result<()> {
             // Override the error severity with warning
             .arg("--ignore")
             .arg("possibly-unresolved-reference"),
-        @r"
+        @r###"
-        success: false
+    success: false
-        exit_code: 1
+    exit_code: 1
-        ----- stdout -----
+    ----- stdout -----
-        warning[lint:division-by-zero] <temp_dir>/test.py:2:5 Cannot divide object of type `Literal[4]` by zero
+    warning: lint:division-by-zero
+     --> <temp_dir>/test.py:2:5
+      |
+    2 | y = 4 / 0
+      |     ----- Cannot divide object of type `Literal[4]` by zero
+    3 |
+    4 | for a in range(0, y):
+      |
 
-        ----- stderr -----
+
-        "
+    ----- stderr -----
+    "###
     );
 
     Ok(())
@@ -329,14 +436,21 @@ fn configuration_unknown_rules() -> anyhow::Result<()> {
         ("test.py", "print(10)"),
     ])?;
 
-    assert_cmd_snapshot!(case.command(), @r"
+    assert_cmd_snapshot!(case.command(), @r###"
-        success: false
+    success: false
-        exit_code: 1
+    exit_code: 1
-        ----- stdout -----
+    ----- stdout -----
-        warning[unknown-rule] <temp_dir>/pyproject.toml:3:1 Unknown lint rule `division-by-zer`
+    warning: unknown-rule
+     --> <temp_dir>/pyproject.toml:3:1
+      |
+    2 | [tool.knot.rules]
+    3 | division-by-zer = "warn" # incorrect rule name
+      | --------------- Unknown lint rule `division-by-zer`
+      |
 
-        ----- stderr -----
+
-    ");
+    ----- stderr -----
+    "###);
 
     Ok(())
 }
@@ -346,14 +460,15 @@ fn configuration_unknown_rules() -> anyhow::Result<()> {
 fn cli_unknown_rules() -> anyhow::Result<()> {
     let case = TestCase::with_file("test.py", "print(10)")?;
 
-    assert_cmd_snapshot!(case.command().arg("--ignore").arg("division-by-zer"), @r"
+    assert_cmd_snapshot!(case.command().arg("--ignore").arg("division-by-zer"), @r###"
-        success: false
+    success: false
-        exit_code: 1
+    exit_code: 1
-        ----- stdout -----
+    ----- stdout -----
-        warning[unknown-rule] Unknown lint rule `division-by-zer`
+    warning: unknown-rule: Unknown lint rule `division-by-zer`
 
-        ----- stderr -----
+
-    ");
+    ----- stderr -----
+    "###);
 
     Ok(())
 }

crates/red_knot_wasm/tests/api.rs

@@ -19,6 +19,15 @@ fn check() {
 
     assert_eq!(
         result,
-        vec!["error[lint:unresolved-import] /test.py:1:8 Cannot resolve import `random22`"]
+        vec![
+            "\
+error: lint:unresolved-import
+ --> /test.py:1:8
+  |
+1 | import random22
+  |        ^^^^^^^^ Cannot resolve import `random22`
+  |
+",
+        ],
     );
 }

crates/ruff_benchmark/benches/red_knot.rs

@@ -1,5 +1,8 @@
 #![allow(clippy::disallowed_names)]
 
+use std::borrow::Cow;
+use std::ops::Range;
+
 use rayon::ThreadPoolBuilder;
 use red_knot_project::metadata::options::{EnvironmentOptions, Options};
 use red_knot_project::metadata::value::RangedValue;
@@ -8,7 +11,7 @@ use red_knot_project::{Db, ProjectDatabase, ProjectMetadata};
 use red_knot_python_semantic::PythonVersion;
 use ruff_benchmark::criterion::{criterion_group, criterion_main, BatchSize, Criterion};
 use ruff_benchmark::TestFile;
-use ruff_db::diagnostic::Diagnostic;
+use ruff_db::diagnostic::{Diagnostic, DiagnosticId, Severity};
 use ruff_db::files::{system_path_to_file, File};
 use ruff_db::source::source_text;
 use ruff_db::system::{MemoryFileSystem, SystemPath, SystemPathBuf, TestSystem};
@@ -23,14 +26,58 @@ struct Case {
 
 const TOMLLIB_312_URL: &str = "https://raw.githubusercontent.com/python/cpython/8e8a4baf652f6e1cee7acde9d78c4b6154539748/Lib/tomllib";
 
-static EXPECTED_DIAGNOSTICS: &[&str] = &[
+/// A structured set of fields we use to do diagnostic comparisons.
+///
+/// This helps assert benchmark results. Previously, we would compare
+/// the actual diagnostic output, but using `insta` inside benchmarks is
+/// problematic, and updating the strings otherwise when diagnostic rendering
+/// changes is a PITA.
+type KeyDiagnosticFields = (
+    DiagnosticId,
+    Option<&'static str>,
+    Option<Range<usize>>,
+    Cow<'static, str>,
+    Severity,
+);
+
+static EXPECTED_DIAGNOSTICS: &[KeyDiagnosticFields] = &[
     // We don't support `*` imports yet:
-    "error[lint:unresolved-import] /src/tomllib/_parser.py:7:29 Module `collections.abc` has no member `Iterable`",
+    (
+        DiagnosticId::lint("unresolved-import"),
+        Some("/src/tomllib/_parser.py"),
+        Some(192..200),
+        Cow::Borrowed("Module `collections.abc` has no member `Iterable`"),
+        Severity::Error,
+    ),
     // We don't handle intersections in `is_assignable_to` yet
-    "error[lint:invalid-argument-type] /src/tomllib/_parser.py:626:46 Object of type `Unknown & ~AlwaysFalsy | @Todo & ~AlwaysFalsy` cannot be assigned to parameter 1 (`match`) of function `match_to_datetime`; expected type `Match`",
+    (
-    "error[lint:invalid-argument-type] /src/tomllib/_parser.py:632:58 Object of type `Unknown & ~AlwaysFalsy | @Todo & ~AlwaysFalsy` cannot be assigned to parameter 1 (`match`) of function `match_to_localtime`; expected type `Match`",
+        DiagnosticId::lint("invalid-argument-type"),
-    "error[lint:invalid-argument-type] /src/tomllib/_parser.py:639:52 Object of type `Unknown & ~AlwaysFalsy | @Todo & ~AlwaysFalsy` cannot be assigned to parameter 1 (`match`) of function `match_to_number`; expected type `Match`",
+        Some("/src/tomllib/_parser.py"),
-    "warning[lint:unused-ignore-comment] /src/tomllib/_parser.py:682:31 Unused blanket `type: ignore` directive",
+        Some(20158..20172),
+        Cow::Borrowed("Object of type `Unknown & ~AlwaysFalsy | @Todo & ~AlwaysFalsy` cannot be assigned to parameter 1 (`match`) of function `match_to_datetime`; expected type `Match`"),
+        Severity::Error,
+    ),
+    (
+        DiagnosticId::lint("invalid-argument-type"),
+        Some("/src/tomllib/_parser.py"),
+        Some(20464..20479),
+        Cow::Borrowed("Object of type `Unknown & ~AlwaysFalsy | @Todo & ~AlwaysFalsy` cannot be assigned to parameter 1 (`match`) of function `match_to_localtime`; expected type `Match`"),
+        Severity::Error,
+    ),
+    (
+        DiagnosticId::lint("invalid-argument-type"),
+        Some("/src/tomllib/_parser.py"),
+        Some(20774..20786),
+        Cow::Borrowed("Object of type `Unknown & ~AlwaysFalsy | @Todo & ~AlwaysFalsy` cannot be assigned to parameter 1 (`match`) of function `match_to_number`; expected type `Match`"),
+        Severity::Error,
+    ),
+    (
+        DiagnosticId::lint("unused-ignore-comment"),
+        Some("/src/tomllib/_parser.py"),
+        Some(22299..22333),
+        Cow::Borrowed("Unused blanket `type: ignore` directive"),
+        Severity::Warning,
+    ),
 ];
 
 fn get_test_file(name: &str) -> TestFile {
@@ -106,7 +153,7 @@ fn benchmark_incremental(criterion: &mut Criterion) {
 
         let result: Vec<_> = case.db.check().unwrap();
 
-        assert_diagnostics(&case.db, result);
+        assert_diagnostics(&case.db, &result);
 
         case.fs
             .write_file(
@@ -151,7 +198,7 @@ fn benchmark_cold(criterion: &mut Criterion) {
                 let Case { db, .. } = case;
                 let result: Vec<_> = db.check().unwrap();
 
-                assert_diagnostics(db, result);
+                assert_diagnostics(db, &result);
             },
             BatchSize::SmallInput,
         );
@@ -159,17 +206,19 @@ fn benchmark_cold(criterion: &mut Criterion) {
 }
 
 #[track_caller]
-fn assert_diagnostics(db: &dyn Db, diagnostics: Vec<Box<dyn Diagnostic>>) {
+fn assert_diagnostics(db: &dyn Db, diagnostics: &[Box<dyn Diagnostic>]) {
     let normalized: Vec<_> = diagnostics
-        .into_iter()
+        .iter()
         .map(|diagnostic| {
-            diagnostic
+            (
-                .display(db.upcast())
+                diagnostic.id(),
-                .to_string()
+                diagnostic.file().map(|file| file.path(db).as_str()),
-                .replace('\\', "/")
+                diagnostic.range().map(Range::<usize>::from),
+                diagnostic.message(),
+                diagnostic.severity(),
+            )
         })
         .collect();
-
     assert_eq!(&normalized, EXPECTED_DIAGNOSTICS);
 }
 

crates/ruff_db/Cargo.toml

@@ -11,6 +11,7 @@ repository = { workspace = true }
 license = { workspace = true }
 
 [dependencies]
+ruff_annotate_snippets = { workspace = true }
 ruff_cache = { workspace = true, optional = true }
 ruff_notebook = { workspace = true }
 ruff_python_ast = { workspace = true }
@@ -20,6 +21,7 @@ ruff_source_file = { workspace = true }
 ruff_text_size = { workspace = true }
 
 camino = { workspace = true }
+colored = { workspace = true }
 countme = { workspace = true }
 dashmap = { workspace = true }
 dunce = { workspace = true }

crates/ruff_db/src/diagnostic.rs

@@ -3,7 +3,9 @@ use std::fmt::Formatter;
 
 use thiserror::Error;
 
+use ruff_annotate_snippets::{Level, Renderer, Snippet};
 use ruff_python_parser::ParseError;
+use ruff_source_file::{OneIndexed, SourceCode};
 use ruff_text_size::TextRange;
 
 use crate::{
@@ -210,29 +212,94 @@ impl<'db> DisplayDiagnostic<'db> {
 
 impl std::fmt::Display for DisplayDiagnostic<'_> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self.diagnostic.severity() {
+        let level = match self.diagnostic.severity() {
-            Severity::Info => f.write_str("info")?,
+            Severity::Info => Level::Info,
-            Severity::Warning => f.write_str("warning")?,
+            Severity::Warning => Level::Warning,
-            Severity::Error => f.write_str("error")?,
+            Severity::Error => Level::Error,
-            Severity::Fatal => f.write_str("fatal")?,
+            // NOTE: Should we really collapse this to "error"?
+            //
+            // After collapsing this, the snapshot tests seem to reveal that we
+            // don't currently have any *tests* with a `fatal` severity level.
+            // And maybe *rendering* this as just an `error` is fine. If we
+            // really do need different rendering, then I think we can add a
+            // `Level::Fatal`. ---AG
+            Severity::Fatal => Level::Error,
+        };
+
+        let render = |f: &mut std::fmt::Formatter, message| {
+            let renderer = if !cfg!(test) && colored::control::SHOULD_COLORIZE.should_colorize() {
+                Renderer::styled()
+            } else {
+                Renderer::plain()
+            }
+            .cut_indicator("…");
+            let rendered = renderer.render(message);
+            writeln!(f, "{rendered}")
+        };
+        match (self.diagnostic.file(), self.diagnostic.range()) {
+            (None, _) => {
+                // NOTE: This is pretty sub-optimal. It doesn't render well. We
+                // really want a snippet, but without a `File`, we can't really
+                // render anything. It looks like this case currently happens
+                // for configuration errors. It looks like we can probably
+                // produce a snippet for this if it comes from a file, but if
+                // it comes from the CLI, I'm not quite sure exactly what to
+                // do. ---AG
+                let msg = format!("{}: {}", self.diagnostic.id(), self.diagnostic.message());
+                render(f, level.title(&msg))
+            }
+            (Some(file), range) => {
+                let path = file.path(self.db).to_string();
+                let source = source_text(self.db, file);
+                let title = self.diagnostic.id().to_string();
+                let message = self.diagnostic.message();
+
+                let Some(range) = range else {
+                    let snippet = Snippet::source(source.as_str()).origin(&path).line_start(1);
+                    return render(f, level.title(&title).snippet(snippet));
+                };
+
+                // The bits below are a simplified copy from
+                // `crates/ruff_linter/src/message/text.rs`.
+                let index = line_index(self.db, file);
+                let source_code = SourceCode::new(source.as_str(), &index);
+
+                let content_start_index = source_code.line_index(range.start());
+                let mut start_index = content_start_index.saturating_sub(2);
+                // Trim leading empty lines.
+                while start_index < content_start_index {
+                    if !source_code.line_text(start_index).trim().is_empty() {
+                        break;
+                    }
+                    start_index = start_index.saturating_add(1);
+                }
+
+                let content_end_index = source_code.line_index(range.end());
+                let mut end_index = content_end_index
+                    .saturating_add(2)
+                    .min(OneIndexed::from_zero_indexed(index.line_count()));
+                // 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);
+                }
+
+                // Slice up the code frame and adjust our range.
+                let start_offset = source_code.line_start(start_index);
+                let end_offset = source_code.line_end(end_index);
+                let frame = source_code.slice(TextRange::new(start_offset, end_offset));
+                let span = range - start_offset;
+
+                let annotation = level.span(span.into()).label(&message);
+                let snippet = Snippet::source(frame)
+                    .origin(&path)
+                    .line_start(start_index.get())
+                    .annotation(annotation);
+                render(f, level.title(&title).snippet(snippet))
+            }
         }
-
-        write!(f, "[{rule}]", rule = self.diagnostic.id())?;
-
-        if let Some(file) = self.diagnostic.file() {
-            write!(f, " {path}", path = file.path(self.db))?;
-        }
-
-        if let (Some(file), Some(range)) = (self.diagnostic.file(), self.diagnostic.range()) {
-            let index = line_index(self.db, file);
-            let source = source_text(self.db, file);
-
-            let start = index.source_location(range.start(), &source);
-
-            write!(f, ":{line}:{col}", line = start.row, col = start.column)?;
-        }
-
-        write!(f, " {message}", message = self.diagnostic.message())
     }
 }