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

Author: MichaReiser

.config/nextest.toml

@@ -7,10 +7,6 @@ serial = { max-threads = 1 }
 filter = 'binary(file_watching)'
 test-group = 'serial'
 
-[[profile.default.overrides]]
-filter = 'binary(e2e)'
-test-group = 'serial'
-
 [profile.ci]
 # Print out output for failing tests as soon as they fail, and also at the end
 # of the run (for easy scrollability).

crates/ruff_diagnostics/src/fix.rs

@@ -149,6 +149,10 @@ impl Fix {
         &self.edits
     }
 
+    pub fn into_edits(self) -> Vec<Edit> {
+        self.edits
+    }
+
     /// Return the [`Applicability`] of the [`Fix`].
     pub fn applicability(&self) -> Applicability {
         self.applicability

crates/ty/docs/rules.md

@@ -1,8 +1,10 @@
 use crate::{completion, find_node::covering_node};
+
 use ruff_db::{files::File, parsed::parsed_module};
 use ruff_diagnostics::Edit;
 use ruff_text_size::TextRange;
 use ty_project::Db;
+use ty_python_semantic::create_suppression_fix;
 use ty_python_semantic::types::UNRESOLVED_REFERENCE;
 
 /// A `QuickFix` Code Action
@@ -18,26 +20,45 @@ pub fn code_actions(
     file: File,
     diagnostic_range: TextRange,
     diagnostic_id: &str,
-) -> Option<Vec<QuickFix>> {
+) -> Vec<QuickFix> {
     let registry = db.lint_registry();
     let Ok(lint_id) = registry.get(diagnostic_id) else {
-        return None;
+        return Vec::new();
     };
-    if lint_id.name() == UNRESOLVED_REFERENCE.name() {
-        let parsed = parsed_module(db, file).load(db);
-        let node = covering_node(parsed.syntax().into(), diagnostic_range).node();
-        let symbol = &node.expr_name()?.id;
 
-        let fixes = completion::missing_imports(db, file, &parsed, symbol, node)
+    let mut actions = Vec::new();
+
+    if lint_id.name() == UNRESOLVED_REFERENCE.name()
+        && let Some(import_quick_fix) = create_import_symbol_quick_fix(db, file, diagnostic_range)
+    {
+        actions.extend(import_quick_fix)
+    }
+
+    actions.push(QuickFix {
+        title: format!("Ignore '{}' for this line", lint_id.name()),
+        edits: create_suppression_fix(db, file, lint_id, diagnostic_range).into_edits(),
+        preferred: false,
+    });
+
+    actions
+}
+
+fn create_import_symbol_quick_fix(
+    db: &dyn Db,
+    file: File,
+    diagnostic_range: TextRange,
+) -> Option<impl Iterator<Item = QuickFix>> {
+    let parsed = parsed_module(db, file).load(db);
+    let node = covering_node(parsed.syntax().into(), diagnostic_range).node();
+    let symbol = &node.expr_name()?.id;
+
+    Some(
+        completion::missing_imports(db, file, &parsed, symbol, node)
             .into_iter()
             .map(|import| QuickFix {
                 title: import.label,
                 edits: vec![import.edit],
                 preferred: true,
-            })
-            .collect();
-        Some(fixes)
-    } else {
-        None
-    }
+            }),
+    )
 }

crates/ty_ide/src/code_action.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/lib.rs

@@ -17,6 +17,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};
 
 declare_lint! {
@@ -375,6 +376,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_python_semantic/src/suppression.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/diagnostics.rs

@@ -82,9 +82,8 @@ impl BackgroundDocumentRequestHandler for CodeActionRequestHandler {
             let encoding = snapshot.encoding();
             if let Some(NumberOrString::String(diagnostic_id)) = &diagnostic.code
                 && let Some(range) = diagnostic.range.to_text_range(db, file, url, encoding)
-                && let Some(fixes) = code_actions(db, file, range, diagnostic_id)
             {
-                for action in fixes {
+                for action in code_actions(db, file, range, diagnostic_id) {
                     actions.push(CodeActionOrCommand::CodeAction(lsp_types::CodeAction {
                         title: action.title,
                         kind: Some(CodeActionKind::QUICKFIX),

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

@@ -51,5 +51,54 @@ expression: code_actions
       }
     },
     "isPreferred": true
+  },
+  {
+    "title": "Ignore 'unused-ignore-comment' for this line",
+    "kind": "quickfix",
+    "diagnostics": [
+      {
+        "range": {
+          "start": {
+            "line": 0,
+            "character": 12
+          },
+          "end": {
+            "line": 0,
+            "character": 42
+          }
+        },
+        "severity": 2,
+        "code": "unused-ignore-comment",
+        "codeDescription": {
+          "href": "https://ty.dev/rules#unused-ignore-comment"
+        },
+        "source": "ty",
+        "message": "Unused `ty: ignore` directive",
+        "relatedInformation": [],
+        "tags": [
+          1
+        ]
+      }
+    ],
+    "edit": {
+      "changes": {
+        "file://<temp_dir>/src/foo.py": [
+          {
+            "range": {
+              "start": {
+                "line": 0,
+                "character": 41
+              },
+              "end": {
+                "line": 0,
+                "character": 41
+              }
+            },
+            "newText": ", unused-ignore-comment"
+          }
+        ]
+      }
+    },
+    "isPreferred": false
   }
 ]