[ty] Add code action support to playground (#21655)

Cargo.lock

@@ -4602,6 +4602,7 @@ dependencies = [
  "js-sys",
  "log",
  "ruff_db",
+ "ruff_diagnostics",
  "ruff_notebook",
  "ruff_python_formatter",
  "ruff_source_file",

crates/ty_wasm/Cargo.toml

@@ -27,6 +27,7 @@ ty_project = { workspace = true, default-features = false, features = [
 ty_python_semantic = { workspace = true }
 
 ruff_db = { workspace = true, default-features = false, features = [] }
+ruff_diagnostics = { workspace = true }
 ruff_notebook = { workspace = true }
 ruff_python_formatter = { workspace = true }
 ruff_source_file = { workspace = true }

crates/ty_wasm/src/lib.rs

@@ -11,6 +11,7 @@ use ruff_db::system::{
     SystemPath, SystemPathBuf, SystemVirtualPath, WritableSystem,
 };
 use ruff_db::vendored::VendoredPath;
+use ruff_diagnostics::Applicability;
 use ruff_notebook::Notebook;
 use ruff_python_formatter::formatted_file;
 use ruff_source_file::{LineIndex, OneIndexed, SourceLocation};
@@ -732,6 +733,53 @@ impl Diagnostic {
             .to_string()
             .into()
     }
+
+    /// Returns the code action for this diagnostic, if it has a fix.
+    #[wasm_bindgen(js_name = "codeAction")]
+    pub fn code_action(&self, workspace: &Workspace) -> Option<CodeAction> {
+        let fix = self
+            .inner
+            .fix()
+            .filter(|fix| fix.applies(Applicability::Unsafe))?;
+
+        let primary_span = self.inner.primary_span()?;
+        let file = primary_span.expect_ty_file();
+
+        let source = source_text(&workspace.db, file);
+        let index = line_index(&workspace.db, file);
+
+        let edits: Vec<TextEdit> = fix
+            .edits()
+            .iter()
+            .map(|edit| TextEdit {
+                range: Range::from_text_range(
+                    edit.range(),
+                    &index,
+                    &source,
+                    workspace.position_encoding,
+                ),
+                new_text: edit.content().unwrap_or_default().to_string(),
+            })
+            .collect();
+
+        let title = self
+            .inner
+            .first_help_text()
+            .map(ToString::to_string)
+            .unwrap_or_else(|| format!("Fix {}", self.inner.id()));
+
+        Some(CodeAction { title, edits })
+    }
+}
+
+/// A code action that can be applied to fix a diagnostic.
+#[wasm_bindgen]
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct CodeAction {
+    #[wasm_bindgen(getter_with_clone)]
+    pub title: String,
+    #[wasm_bindgen(getter_with_clone)]
+    pub edits: Vec<TextEdit>,
 }
 
 #[wasm_bindgen]

playground/ty/src/Editor/Chrome.tsx

@@ -355,6 +355,7 @@ function useCheckResult(
         severity: diagnostic.severity(),
         range: diagnostic.toRange(workspace) ?? null,
         textRange: diagnostic.textRange() ?? null,
+        raw: diagnostic,
       }));
 
       return {

playground/ty/src/Editor/Diagnostics.tsx

@@ -1,4 +1,9 @@
-import type { Severity, Range, TextRange } from "ty_wasm";
+import type {
+  Severity,
+  Range,
+  TextRange,
+  Diagnostic as TyDiagnostic,
+} from "ty_wasm";
 import classNames from "classnames";
 import { Theme } from "shared";
 import { useMemo } from "react";
@@ -103,4 +108,5 @@ export interface Diagnostic {
   severity: Severity;
   range: Range | null;
   textRange: TextRange | null;
+  raw: TyDiagnostic;
 }

playground/ty/src/Editor/Editor.tsx

@@ -189,8 +189,11 @@ class PlaygroundServer
     languages.DocumentSemanticTokensProvider,
     languages.DocumentRangeSemanticTokensProvider,
     languages.SignatureHelpProvider,
-    languages.DocumentHighlightProvider
+    languages.DocumentHighlightProvider,
+    languages.CodeActionProvider
 {
+  private diagnostics: Diagnostic[] = [];
+
   private typeDefinitionProviderDisposable: IDisposable;
   private declarationProviderDisposable: IDisposable;
   private definitionProviderDisposable: IDisposable;
@@ -204,6 +207,7 @@ class PlaygroundServer
   private rangeSemanticTokensDisposable: IDisposable;
   private signatureHelpDisposable: IDisposable;
   private documentHighlightDisposable: IDisposable;
+  private codeActionDisposable: IDisposable;
   private inVendoredFileCondition: editor.IContextKey<boolean>;
   // Cache for vendored file handles
   private vendoredFileHandles = new Map<string, FileHandle>();
@@ -253,6 +257,10 @@ class PlaygroundServer
       monaco.languages.registerSignatureHelpProvider("python", this);
     this.documentHighlightDisposable =
       monaco.languages.registerDocumentHighlightProvider("python", this);
+    this.codeActionDisposable = monaco.languages.registerCodeActionProvider(
+      "python",
+      this,
+    );
 
     this.inVendoredFileCondition = editor.createContextKey<boolean>(
       "inVendoredFile",
@@ -534,6 +542,8 @@ class PlaygroundServer
   }
 
   updateDiagnostics(diagnostics: Array<Diagnostic>) {
+    this.diagnostics = diagnostics;
+
     if (this.props.files.selected == null) {
       return;
     }
@@ -584,6 +594,59 @@ class PlaygroundServer
     );
   }
 
+  provideCodeActions(
+    model: editor.ITextModel,
+    range: Range,
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    _context: languages.CodeActionContext,
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    _token: CancellationToken,
+  ): languages.ProviderResult<languages.CodeActionList> {
+    const actions: languages.CodeAction[] = [];
+
+    for (const diagnostic of this.diagnostics) {
+      const diagnosticRange = diagnostic.range;
+      if (diagnosticRange == null) {
+        continue;
+      }
+
+      const monacoRange = tyRangeToMonacoRange(diagnosticRange);
+      if (!Range.areIntersecting(range, monacoRange)) {
+        continue;
+      }
+
+      const codeAction = diagnostic.raw.codeAction(this.props.workspace);
+      if (codeAction == null) {
+        continue;
+      }
+
+      actions.push({
+        title: codeAction.title,
+        kind: "quickfix",
+        isPreferred: true,
+        edit: {
+          edits: codeAction.edits.map((edit) => ({
+            resource: model.uri,
+            textEdit: {
+              range: tyRangeToMonacoRange(edit.range),
+              text: edit.new_text,
+            },
+            versionId: model.getVersionId(),
+          })),
+        },
+      });
+    }
+
+    if (actions.length === 0) {
+      return undefined;
+    }
+
+    return {
+      actions,
+      dispose: () => {},
+    };
+  }
+
   provideHover(
     model: editor.ITextModel,
     position: Position,
@@ -836,6 +899,7 @@ class PlaygroundServer
     this.completionDisposable.dispose();
     this.signatureHelpDisposable.dispose();
     this.documentHighlightDisposable.dispose();
+    this.codeActionDisposable.dispose();
   }
 }