Use Annotation::tags instead of hardcoded rule matching in ruff server (#20565)

Co-authored-by: Micha Reiser <micha@reiser.io>

crates/ruff_linter/src/checkers/ast/mod.rs

@@ -28,7 +28,7 @@ use itertools::Itertools;
 use log::debug;
 use rustc_hash::{FxHashMap, FxHashSet};
 
-use ruff_db::diagnostic::{Annotation, Diagnostic, IntoDiagnosticMessage, Span};
+use ruff_db::diagnostic::{Annotation, Diagnostic, DiagnosticTag, IntoDiagnosticMessage, Span};
 use ruff_diagnostics::{Applicability, Fix, IsolationLevel};
 use ruff_notebook::{CellOffsets, NotebookIndex};
 use ruff_python_ast::helpers::{collect_import_from_member, is_docstring_stmt, to_module_path};
@@ -3326,6 +3326,56 @@ impl DiagnosticGuard<'_, '_> {
     pub(crate) fn defuse(mut self) {
         self.diagnostic = None;
     }
+
+    /// Set the message on the primary annotation for this diagnostic.
+    ///
+    /// If a message already exists on the primary annotation, then this
+    /// overwrites the existing message.
+    ///
+    /// This message is associated with the primary annotation created
+    /// for every `Diagnostic` that uses the `DiagnosticGuard` API.
+    /// Specifically, the annotation is derived from the `TextRange` given to
+    /// the `LintContext::report_diagnostic` API.
+    ///
+    /// Callers can add additional primary or secondary annotations via the
+    /// `DerefMut` trait implementation to a `Diagnostic`.
+    pub(crate) fn set_primary_message(&mut self, message: impl IntoDiagnosticMessage) {
+        // N.B. It is normally bad juju to define `self` methods
+        // on types that implement `Deref`. Instead, it's idiomatic
+        // to do `fn foo(this: &mut LintDiagnosticGuard)`, which in
+        // turn forces callers to use
+        // `LintDiagnosticGuard(&mut guard, message)`. But this is
+        // supremely annoying for what is expected to be a common
+        // case.
+        //
+        // Moreover, most of the downside that comes from these sorts
+        // of methods is a semver hazard. Because the deref target type
+        // could also define a method by the same name, and that leads
+        // to confusion. But we own all the code involved here and
+        // there is no semver boundary. So... ¯\_(ツ)_/¯ ---AG
+
+        // OK because we know the diagnostic was constructed with a single
+        // primary annotation that will always come before any other annotation
+        // in the diagnostic. (This relies on the `Diagnostic` API not exposing
+        // any methods for removing annotations or re-ordering them, which is
+        // true as of 2025-04-11.)
+        let ann = self.primary_annotation_mut().unwrap();
+        ann.set_message(message);
+    }
+
+    /// Adds a tag on the primary annotation for this diagnostic.
+    ///
+    /// This tag is associated with the primary annotation created
+    /// for every `Diagnostic` that uses the `DiagnosticGuard` API.
+    /// Specifically, the annotation is derived from the `TextRange` given to
+    /// the `LintContext::report_diagnostic` API.
+    ///
+    /// Callers can add additional primary or secondary annotations via the
+    /// `DerefMut` trait implementation to a `Diagnostic`.
+    pub(crate) fn add_primary_tag(&mut self, tag: DiagnosticTag) {
+        let ann = self.primary_annotation_mut().unwrap();
+        ann.push_tag(tag);
+    }
 }
 
 impl DiagnosticGuard<'_, '_> {

crates/ruff_linter/src/rules/pyflakes/rules/redefined_while_unused.rs

@@ -197,9 +197,7 @@ pub(crate) fn redefined_while_unused(checker: &Checker, scope_id: ScopeId, scope
                 shadowed,
             );
 
-            if let Some(ann) = diagnostic.primary_annotation_mut() {
+            diagnostic.set_primary_message(format_args!("`{name}` redefined here"));
-                ann.set_message(format_args!("`{name}` redefined here"));
-            }
 
             if let Some(range) = binding.parent_range(checker.semantic()) {
                 diagnostic.set_parent(range.start());

crates/ruff_linter/src/rules/pyflakes/rules/unused_import.rs

@@ -435,6 +435,8 @@ pub(crate) fn unused_import(checker: &Checker, scope: &Scope) {
                     diagnostic.set_fix(fix.clone());
                 }
             }
+
+            diagnostic.add_primary_tag(ruff_db::diagnostic::DiagnosticTag::Unnecessary);
         }
     }
 
@@ -455,6 +457,8 @@ pub(crate) fn unused_import(checker: &Checker, scope: &Scope) {
         if let Some(range) = binding.parent_range {
             diagnostic.set_parent(range.start());
         }
+
+        diagnostic.add_primary_tag(ruff_db::diagnostic::DiagnosticTag::Unnecessary);
     }
 }
 

crates/ruff_linter/src/rules/pyflakes/rules/unused_variable.rs

@@ -272,4 +272,6 @@ pub(crate) fn unused_variable(checker: &Checker, name: &str, binding: &Binding)
     if let Some(fix) = remove_unused_variable(binding, checker) {
         diagnostic.set_fix(fix);
     }
+    // Add Unnecessary tag for unused variables
+    diagnostic.add_primary_tag(ruff_db::diagnostic::DiagnosticTag::Unnecessary);
 }

crates/ruff_linter/src/rules/pyupgrade/rules/deprecated_import.rs

@@ -755,6 +755,7 @@ pub(crate) fn deprecated_import(checker: &Checker, import_from_stmt: &StmtImport
             },
             import_from_stmt.range(),
         );
+        diagnostic.add_primary_tag(ruff_db::diagnostic::DiagnosticTag::Deprecated);
         if let Some(content) = fix {
             diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
                 content,

crates/ruff_linter/src/rules/ruff/rules/unused_unpacked_variable.rs

@@ -100,4 +100,6 @@ pub(crate) fn unused_unpacked_variable(checker: &Checker, name: &str, binding: &
     if let Some(fix) = remove_unused_variable(binding, checker) {
         diagnostic.set_fix(fix);
     }
+    // Add Unnecessary tag for unused unpacked variables
+    diagnostic.add_primary_tag(ruff_db::diagnostic::DiagnosticTag::Unnecessary);
 }

crates/ruff_server/src/lint.rs

@@ -287,7 +287,7 @@ fn to_lsp_diagnostic(
         let code = code.to_string();
         (
             Some(severity(&code)),
-            tags(&code),
+            tags(diagnostic),
             Some(lsp_types::NumberOrString::String(code)),
         )
     } else {
@@ -338,12 +338,17 @@ fn severity(code: &str) -> lsp_types::DiagnosticSeverity {
     }
 }
 
-fn tags(code: &str) -> Option<Vec<lsp_types::DiagnosticTag>> {
+fn tags(diagnostic: &Diagnostic) -> Option<Vec<lsp_types::DiagnosticTag>> {
-    match code {
+    diagnostic.primary_tags().map(|tags| {
-        // F401: <module> imported but unused
+        tags.iter()
-        // F841: local variable <name> is assigned to but never used
+            .map(|tag| match tag {
-        // RUF059: Unused unpacked variable
+                ruff_db::diagnostic::DiagnosticTag::Unnecessary => {
-        "F401" | "F841" | "RUF059" => Some(vec![lsp_types::DiagnosticTag::UNNECESSARY]),
+                    lsp_types::DiagnosticTag::UNNECESSARY
-        _ => None,
                 }
+                ruff_db::diagnostic::DiagnosticTag::Deprecated => {
+                    lsp_types::DiagnosticTag::DEPRECATED
+                }
+            })
+            .collect()
+    })
 }