From b79d43a85260c261ca014e3d6e70658be5285cd3 Mon Sep 17 00:00:00 2001
From: Andrew Gallant <andrew@astral.sh>
Date: Fri, 11 Apr 2025 10:16:31 -0400
Subject: [PATCH] ruff_db: add primary annotation message mutators on
 `Diagnostic`

This will enable us to provide an API on `LintDiagnosticGuard` for
setting the primary annotation message. It will require an `unwrap()`,
but due to how `LintDiagnosticGuard` will build a `Diagnostic`, this
`unwrap()` will be guaranteed to succeed. (And it won't bubble out to
every user of `LintDiagnosticGuard`.)
---
 crates/ruff_db/src/diagnostic/mod.rs | 30 +++++++++++++++++++++++++---
 1 file changed, 27 insertions(+), 3 deletions(-)

diff --git a/crates/ruff_db/src/diagnostic/mod.rs b/crates/ruff_db/src/diagnostic/mod.rs
index c638dde8e1..1e2d1785b3 100644
--- a/crates/ruff_db/src/diagnostic/mod.rs
+++ b/crates/ruff_db/src/diagnostic/mod.rs
@@ -198,14 +198,27 @@ impl Diagnostic {
         self.inner.severity
     }
 
-    /// Returns the "primary" annotation of this diagnostic if one exists.
+    /// Returns a shared borrow of the "primary" annotation of this diagnostic
+    /// if one exists.
     ///
-    /// When there are multiple primary annotation, then the first one that was
-    /// added to this diagnostic is returned.
+    /// When there are multiple primary annotations, then the first one that
+    /// was added to this diagnostic is returned.
     pub fn primary_annotation(&self) -> Option<&Annotation> {
         self.inner.annotations.iter().find(|ann| ann.is_primary)
     }
 
+    /// Returns a mutable borrow of the "primary" annotation of this diagnostic
+    /// if one exists.
+    ///
+    /// When there are multiple primary annotations, then the first one that
+    /// was added to this diagnostic is returned.
+    pub fn primary_annotation_mut(&mut self) -> Option<&mut Annotation> {
+        Arc::make_mut(&mut self.inner)
+            .annotations
+            .iter_mut()
+            .find(|ann| ann.is_primary)
+    }
+
     /// Returns the "primary" span of this diagnostic if one exists.
     ///
     /// When there are multiple primary spans, then the first one that was
@@ -378,6 +391,17 @@ impl Annotation {
         Annotation { message, ..self }
     }
 
+    /// Sets the message on this annotation.
+    ///
+    /// If one was already set, then this overwrites it.
+    ///
+    /// This is useful if one needs to set the message on an annotation,
+    /// and all one has is a `&mut Annotation`. For example, via
+    /// `Diagnostic::primary_annotation_mut`.
+    pub fn set_message<'a>(&mut self, message: impl IntoDiagnosticMessage + 'a) {
+        self.message = Some(message.into_diagnostic_message());
+    }
+
     /// Returns the message attached to this annotation, if one exists.
     pub fn get_message(&self) -> Option<&str> {
         self.message.as_ref().map(|m| m.as_str())