[ty] Add code action to suppress diagnostic on current line

Author: MichaReiser

crates/ty_python_semantic/src/lib.rs

@@ -25,6 +25,7 @@ pub use semantic_model::{
     Completion, HasDefinition, HasType, MemberDefinition, NameKind, SemanticModel,
 };
 pub use site_packages::{PythonEnvironment, SitePackagesPaths, SysPrefixPathOrigin};
+pub use suppression::create_suppression_fix;
 pub use types::DisplaySettings;
 pub use types::ide_support::{
     ImportAliasResolution, ResolvedDefinition, definitions_for_attribute, definitions_for_bin_op,

crates/ty_python_semantic/src/suppression.rs

@@ -10,6 +10,7 @@ use ruff_db::{files::File, parsed::parsed_module, source::source_text};
 use ruff_diagnostics::{Edit, Fix};
 use ruff_python_parser::TokenKind;
 use ruff_python_trivia::Cursor;
+use ruff_source_file::LineRanges;
 use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
 use smallvec::{SmallVec, smallvec};
 use std::error::Error;
@@ -384,6 +385,60 @@ fn check_unused_suppressions(context: &mut CheckSuppressionsContext) {
     }
 }
 
+/// Creates a fix for adding a suppression comment to suppress `lint` for `range`.
+///
+/// The fix prefers adding the code to an existing `ty: ignore[]` comment over
+/// adding a new suppression comment.
+pub fn create_suppression_fix(db: &dyn Db, file: File, id: LintId, range: TextRange) -> Fix {
+    let suppressions = suppressions(db, file);
+    let source = source_text(db, file);
+
+    let mut existing_suppressions = suppressions.line_suppressions(range).filter(|suppression| {
+        matches!(
+            suppression.target,
+            SuppressionTarget::Lint(_) | SuppressionTarget::Empty,
+        )
+    });
+
+    // If there's an existing `ty: ignore[]` comment, append the code to it instead of creating a new suppression comment.
+    if let Some(existing) = existing_suppressions.next() {
+        let comment_text = &source[existing.comment_range];
+        if let Some(before_closing_paren) = comment_text.trim_end().strip_suffix(']') {
+            let up_to_last_code = before_closing_paren.trim_end();
+
+            let insertion = if up_to_last_code.ends_with(',') {
+                format!(" {id}", id = id.name())
+            } else {
+                format!(", {id}", id = id.name())
+            };
+
+            let relative_offset_from_end = comment_text.text_len() - up_to_last_code.text_len();
+
+            return Fix::safe_edit(Edit::insertion(
+                insertion,
+                existing.comment_range.end() - relative_offset_from_end,
+            ));
+        }
+    }
+
+    // Always insert a new suppression at the end of the range to avoid having to deal with multiline strings
+    // etc.
+    let line_end = source.line_end(range.end());
+    let up_to_line_end = &source[..line_end.to_usize()];
+    let up_to_first_content = up_to_line_end.trim_end();
+    let trailing_whitespace_len = up_to_line_end.text_len() - up_to_first_content.text_len();
+
+    let insertion = format!("  # ty:ignore[{id}]", id = id.name());
+
+    Fix::safe_edit(if trailing_whitespace_len == TextSize::ZERO {
+        Edit::insertion(insertion, line_end)
+    } else {
+        // `expr # fmt: off<trailing_whitespace>`
+        // Trim the trailing whitespace
+        Edit::replacement(insertion, line_end - trailing_whitespace_len, line_end)
+    })
+}
+
 struct CheckSuppressionsContext<'a> {
     db: &'a dyn Db,
     file: File,

crates/ty_server/src/server/api/diagnostics.rs

@@ -6,7 +6,7 @@ use lsp_types::{
     CodeDescription, Diagnostic, DiagnosticRelatedInformation, DiagnosticSeverity, DiagnosticTag,
     NumberOrString, PublishDiagnosticsParams, Url,
 };
-use ruff_diagnostics::Applicability;
+use ruff_diagnostics::{Applicability, Fix};
 use ruff_text_size::Ranged;
 use rustc_hash::FxHashMap;
 
@@ -428,29 +428,38 @@ impl DiagnosticData {
         let primary_span = diagnostic.primary_span()?;
         let file = primary_span.expect_ty_file();
 
-        let mut lsp_edits: HashMap<Url, Vec<lsp_types::TextEdit>> = HashMap::new();
-
-        for edit in fix.edits() {
-            let location = edit
-                .range()
-                .to_lsp_range(db, file, encoding)?
-                .to_location()?;
-
-            lsp_edits
-                .entry(location.uri)
-                .or_default()
-                .push(lsp_types::TextEdit {
-                    range: location.range,
-                    new_text: edit.content().unwrap_or_default().to_string(),
-                });
-        }
-
         Some(Self {
             fix_title: diagnostic
                 .first_help_text()
                 .map(ToString::to_string)
                 .unwrap_or_else(|| format!("Fix {}", diagnostic.id())),
-            edits: lsp_edits,
+            edits: fix_to_lsp_edits(db, fix, file, encoding)?,
         })
     }
 }
+
+pub(super) fn fix_to_lsp_edits(
+    db: &dyn Db,
+    fix: &Fix,
+    file: File,
+    encoding: PositionEncoding,
+) -> Option<HashMap<Url, Vec<lsp_types::TextEdit>>> {
+    let mut lsp_edits: HashMap<Url, Vec<lsp_types::TextEdit>> = HashMap::new();
+
+    for edit in fix.edits() {
+        let location = edit
+            .range()
+            .to_lsp_range(db, file, encoding)?
+            .to_location()?;
+
+        lsp_edits
+            .entry(location.uri)
+            .or_default()
+            .push(lsp_types::TextEdit {
+                range: location.range,
+                new_text: edit.content().unwrap_or_default().to_string(),
+            });
+    }
+
+    Some(lsp_edits)
+}

crates/ty_server/src/server/api/requests/code_action.rs

@@ -1,13 +1,15 @@
 use std::borrow::Cow;
 
-use lsp_types::{self as types, Url, request as req};
+use lsp_types::{self as types, NumberOrString, Url, request as req};
 use ty_project::ProjectDatabase;
+use ty_python_semantic::{Db as _, create_suppression_fix};
 use types::{CodeActionKind, CodeActionOrCommand};
 
 use crate::DIAGNOSTIC_NAME;
+use crate::document::RangeExt;
 use crate::server::Result;
 use crate::server::api::RequestHandler;
-use crate::server::api::diagnostics::DiagnosticData;
+use crate::server::api::diagnostics::{DiagnosticData, fix_to_lsp_edits};
 use crate::server::api::traits::{BackgroundDocumentRequestHandler, RetriableRequestHandler};
 use crate::session::DocumentSnapshot;
 use crate::session::client::Client;
@@ -24,45 +26,89 @@ impl BackgroundDocumentRequestHandler for CodeActionRequestHandler {
     }
 
     fn run_with_snapshot(
-        _db: &ProjectDatabase,
-        _snapshot: &DocumentSnapshot,
+        db: &ProjectDatabase,
+        snapshot: &DocumentSnapshot,
         _client: &Client,
         params: types::CodeActionParams,
     ) -> Result<Option<types::CodeActionResponse>> {
         let diagnostics = params.context.diagnostics;
 
         let mut actions = Vec::new();
+        let lint_registry = db.lint_registry();
+        let file = snapshot.to_notebook_or_file(db);
+
+        tracing::debug!("code actions with diagnostics: {:#?}", diagnostics);
 
         for mut diagnostic in diagnostics.into_iter().filter(|diagnostic| {
             diagnostic.source.as_deref() == Some(DIAGNOSTIC_NAME)
                 && range_intersect(&diagnostic.range, &params.range)
         }) {
-            let Some(data) = diagnostic.data.take() else {
-                continue;
+            tracing::debug!("Lint code: {:?}", diagnostic.code.as_ref());
+
+            let add_suppression_fix = if let Some(NumberOrString::String(code)) =
+                diagnostic.code.as_ref()
+                && let Ok(lint) = lint_registry.get(code)
+                && let Some(file) = file
+                && let Some(range) = diagnostic.range.to_text_range(
+                    db,
+                    file,
+                    &params.text_document.uri,
+                    snapshot.encoding(),
+                ) {
+                tracing::debug!("Creating suppression fix");
+                let fix = create_suppression_fix(db, file, lint, range);
+                tracing::debug!("Suppression fix: {fix:#?}");
+                Some((fix, lint))
+            } else {
+                None
             };
 
-            let data: DiagnosticData = match serde_json::from_value(data) {
-                Ok(data) => data,
-                Err(err) => {
-                    tracing::warn!("Failed to deserialize diagnostic data: {err}");
-                    continue;
-                }
+            // If the diagnostic has a fix, add it as a code action.
+            if let Some(data) = diagnostic.data.take() {
+                let data: DiagnosticData = match serde_json::from_value(data) {
+                    Ok(data) => data,
+                    Err(err) => {
+                        tracing::warn!("Failed to deserialize diagnostic data: {err}");
+                        continue;
+                    }
+                };
+
+                actions.push(CodeActionOrCommand::CodeAction(lsp_types::CodeAction {
+                    title: data.fix_title,
+                    kind: Some(CodeActionKind::QUICKFIX),
+                    diagnostics: Some(vec![diagnostic]),
+                    edit: Some(lsp_types::WorkspaceEdit {
+                        changes: Some(data.edits),
+                        document_changes: None,
+                        change_annotations: None,
+                    }),
+                    is_preferred: Some(true),
+                    command: None,
+                    disabled: None,
+                    data: None,
+                }));
             };
 
-            actions.push(CodeActionOrCommand::CodeAction(lsp_types::CodeAction {
-                title: data.fix_title,
-                kind: Some(CodeActionKind::QUICKFIX),
-                diagnostics: Some(vec![diagnostic]),
-                edit: Some(lsp_types::WorkspaceEdit {
-                    changes: Some(data.edits),
-                    document_changes: None,
-                    change_annotations: None,
-                }),
-                is_preferred: Some(true),
-                command: None,
-                disabled: None,
-                data: None,
-            }));
+            if let Some((suppression_fix, lint)) = add_suppression_fix
+                && let Some(file) = file
+                && let Some(suppression_edits) =
+                    fix_to_lsp_edits(db, &suppression_fix, file, snapshot.encoding())
+            {
+                actions.push(CodeActionOrCommand::CodeAction(lsp_types::CodeAction {
+                    title: format!("Ignore '{}' for this line", lint.name()),
+                    kind: Some(CodeActionKind::QUICKFIX),
+                    diagnostics: None,
+                    edit: Some(lsp_types::WorkspaceEdit {
+                        changes: Some(suppression_edits),
+                        document_changes: None,
+                        change_annotations: None,
+                    }),
+                    is_preferred: Some(false),
+                    command: None,
+                    disabled: None,
+                    data: None,
+                }))
+            }
         }
 
         if actions.is_empty() {