[red-knot] colorize concise output diagnostics (#17232) (#17479)

Co-authored-by: Micha Reiser <micha@reiser.io>
Co-authored-by: Andrew Gallant <andrew@astral.sh>

Cargo.lock

@@ -2865,6 +2865,7 @@ dependencies = [
 name = "ruff_db"
 version = "0.0.0"
 dependencies = [
+ "anstyle",
  "camino",
  "countme",
  "dashmap 6.1.0",

crates/ruff_db/Cargo.toml

@@ -20,6 +20,7 @@ ruff_python_trivia = { workspace = true }
 ruff_source_file = { workspace = true }
 ruff_text_size = { workspace = true }
 
+anstyle = { workspace = true }
 camino = { workspace = true }
 countme = { workspace = true }
 dashmap = { workspace = true }

crates/ruff_db/src/diagnostic/mod.rs

@@ -11,6 +11,7 @@ use crate::Db;
 
 use self::render::FileResolver;
 mod render;
+mod stylesheet;
 
 /// A collection of information that can be rendered into a diagnostic.
 ///

crates/ruff_db/src/diagnostic/render.rs

@@ -7,6 +7,7 @@ use ruff_annotate_snippets::{
 use ruff_source_file::{LineIndex, OneIndexed, SourceCode};
 use ruff_text_size::{TextRange, TextSize};
 
+use crate::diagnostic::stylesheet::{fmt_styled, DiagnosticStylesheet};
 use crate::{
     files::File,
     source::{line_index, source_text, SourceText},
@@ -48,6 +49,7 @@ impl<'a> DisplayDiagnostic<'a> {
         } else {
             AnnotateRenderer::plain()
         };
+
         DisplayDiagnostic {
             config,
             resolver,
@@ -59,31 +61,64 @@ impl<'a> DisplayDiagnostic<'a> {
 
 impl std::fmt::Display for DisplayDiagnostic<'_> {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
-        if matches!(self.config.format, DiagnosticFormat::Concise) {
-            match self.diag.severity() {
-                Severity::Info => f.write_str("info")?,
-                Severity::Warning => f.write_str("warning")?,
-                Severity::Error => f.write_str("error")?,
-                Severity::Fatal => f.write_str("fatal")?,
-            }
+        let stylesheet = if self.config.color {
+            DiagnosticStylesheet::styled()
+        } else {
+            DiagnosticStylesheet::plain()
+        };
+
+        if matches!(self.config.format, DiagnosticFormat::Concise) {
+            let (severity, severity_style) = match self.diag.severity() {
+                Severity::Info => ("info", stylesheet.info),
+                Severity::Warning => ("warning", stylesheet.warning),
+                Severity::Error => ("error", stylesheet.error),
+                Severity::Fatal => ("fatal", stylesheet.error),
+            };
+
+            write!(
+                f,
+                "{severity}[{id}]",
+                severity = fmt_styled(severity, severity_style),
+                id = fmt_styled(self.diag.id(), stylesheet.emphasis)
+            )?;
 
-            write!(f, "[{rule}]", rule = self.diag.id())?;
             if let Some(span) = self.diag.primary_span() {
-                write!(f, " {path}", path = self.resolver.path(span.file()))?;
+                write!(
+                    f,
+                    " {path}",
+                    path = fmt_styled(self.resolver.path(span.file()), stylesheet.emphasis)
+                )?;
                 if let Some(range) = span.range() {
                     let input = self.resolver.input(span.file());
                     let start = input.as_source_code().line_column(range.start());
-                    write!(f, ":{line}:{col}", line = start.line, col = start.column)?;
+
+                    write!(
+                        f,
+                        ":{line}:{col}",
+                        line = fmt_styled(start.line, stylesheet.emphasis),
+                        col = fmt_styled(start.column, stylesheet.emphasis),
+                    )?;
                 }
                 write!(f, ":")?;
             }
-            return writeln!(f, " {}", self.diag.concise_message());
+            return writeln!(f, " {message}", message = self.diag.concise_message());
         }
 
+        let mut renderer = self.annotate_renderer.clone();
+        renderer = renderer
+            .error(stylesheet.error)
+            .warning(stylesheet.warning)
+            .info(stylesheet.info)
+            .note(stylesheet.note)
+            .help(stylesheet.help)
+            .line_no(stylesheet.line_no)
+            .emphasis(stylesheet.emphasis)
+            .none(stylesheet.none);
+
         let resolved = Resolved::new(&self.resolver, self.diag);
         let renderable = resolved.to_renderable(self.config.context);
         for diag in renderable.diagnostics.iter() {
-            writeln!(f, "{}", self.annotate_renderer.render(diag.to_annotate()))?;
+            writeln!(f, "{}", renderer.render(diag.to_annotate()))?;
         }
         writeln!(f)
     }

crates/ruff_db/src/diagnostic/stylesheet.rs

@@ -0,0 +1,80 @@
+use anstyle::{AnsiColor, Effects, Style};
+use std::fmt::Formatter;
+
+pub(super) const fn fmt_styled<'a, T>(
+    content: T,
+    style: anstyle::Style,
+) -> impl std::fmt::Display + 'a
+where
+    T: std::fmt::Display + 'a,
+{
+    struct FmtStyled<T> {
+        content: T,
+        style: anstyle::Style,
+    }
+
+    impl<T> std::fmt::Display for FmtStyled<T>
+    where
+        T: std::fmt::Display,
+    {
+        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+            write!(
+                f,
+                "{style_start}{content}{style_end}",
+                style_start = self.style.render(),
+                content = self.content,
+                style_end = self.style.render_reset()
+            )
+        }
+    }
+
+    FmtStyled { content, style }
+}
+
+#[derive(Clone, Debug)]
+pub struct DiagnosticStylesheet {
+    pub(crate) error: Style,
+    pub(crate) warning: Style,
+    pub(crate) info: Style,
+    pub(crate) note: Style,
+    pub(crate) help: Style,
+    pub(crate) line_no: Style,
+    pub(crate) emphasis: Style,
+    pub(crate) none: Style,
+}
+
+impl Default for DiagnosticStylesheet {
+    fn default() -> Self {
+        Self::plain()
+    }
+}
+
+impl DiagnosticStylesheet {
+    /// Default terminal styling
+    pub fn styled() -> Self {
+        let bright_blue = AnsiColor::BrightBlue.on_default();
+        Self {
+            error: AnsiColor::BrightRed.on_default().effects(Effects::BOLD),
+            warning: AnsiColor::Yellow.on_default().effects(Effects::BOLD),
+            info: bright_blue.effects(Effects::BOLD),
+            note: AnsiColor::BrightGreen.on_default().effects(Effects::BOLD),
+            help: AnsiColor::BrightCyan.on_default().effects(Effects::BOLD),
+            line_no: bright_blue.effects(Effects::BOLD),
+            emphasis: Style::new().effects(Effects::BOLD),
+            none: Style::new(),
+        }
+    }
+
+    pub fn plain() -> Self {
+        Self {
+            error: Style::new(),
+            warning: Style::new(),
+            info: Style::new(),
+            note: Style::new(),
+            help: Style::new(),
+            line_no: Style::new(),
+            emphasis: Style::new(),
+            none: Style::new(),
+        }
+    }
+}