[ty] followup: add-import action for reveal_type too (#21668)

Author: Gankra

crates/ty_ide/src/code_action.rs

@@ -5,7 +5,8 @@ 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;
+use ty_python_semantic::lint::LintId;
+use ty_python_semantic::types::{UNDEFINED_REVEAL, UNRESOLVED_REFERENCE};
 
 /// A `QuickFix` Code Action
 #[derive(Debug, Clone)]
@@ -28,12 +29,17 @@ pub fn code_actions(
 
     let mut actions = Vec::new();
 
-    if lint_id.name() == UNRESOLVED_REFERENCE.name()
+    // Suggest imports for unresolved references (often ideal)
+    // TODO: suggest qualifying with an already imported symbol
+    let is_unresolved_reference =
+        lint_id == LintId::of(&UNRESOLVED_REFERENCE) || lint_id == LintId::of(&UNDEFINED_REVEAL);
+    if is_unresolved_reference
         && let Some(import_quick_fix) = create_import_symbol_quick_fix(db, file, diagnostic_range)
     {
         actions.extend(import_quick_fix);
     }
 
+    // Suggest just suppressing the lint (always a valid option, but never ideal)
     actions.push(QuickFix {
         title: format!("Ignore '{}' for this line", lint_id.name()),
         edits: create_suppression_fix(db, file, lint_id, diagnostic_range).into_edits(),

crates/ty_server/tests/e2e/code_actions.rs

@@ -132,11 +132,44 @@ x: Literal[1] = 1
 ";
 
     let ty_toml = SystemPath::new("ty.toml");
-    let ty_toml_content = "\
-[rules]
-unused-ignore-comment = \"warn\"
+    let ty_toml_content = "";
+
+    let mut server = TestServerBuilder::new()?
+        .with_workspace(workspace_root, None)?
+        .with_file(ty_toml, ty_toml_content)?
+        .with_file(foo, foo_content)?
+        .enable_pull_diagnostics(true)
+        .build()
+        .wait_until_workspaces_are_initialized();
+
+    server.open_text_document(foo, foo_content, 1);
+
+    // Wait for diagnostics to be computed.
+    let diagnostics = server.document_diagnostic_request(foo, None);
+    let range = full_range(foo_content);
+    let code_action_params = code_actions_at(&server, diagnostics, foo, range);
+
+    // Get code actions
+    let code_action_id = server.send_request::<CodeActionRequest>(code_action_params);
+    let code_actions = server.await_response::<CodeActionRequest>(&code_action_id);
+
+    insta::assert_json_snapshot!(code_actions);
+
+    Ok(())
+}
+
+// `Literal` is available from two places so we should suggest two possible imports
+#[test]
+fn code_action_undefined_reveal_type() -> Result<()> {
+    let workspace_root = SystemPath::new("src");
+    let foo = SystemPath::new("src/foo.py");
+    let foo_content = "\
+reveal_type(1)
 ";
 
+    let ty_toml = SystemPath::new("ty.toml");
+    let ty_toml_content = "";
+
     let mut server = TestServerBuilder::new()?
         .with_workspace(workspace_root, None)?
         .with_file(ty_toml, ty_toml_content)?

crates/ty_server/tests/e2e/snapshots/e2e__code_actions__code_action_undefined_reveal_type.snap

@@ -0,0 +1,98 @@
+---
+source: crates/ty_server/tests/e2e/code_actions.rs
+expression: code_actions
+---
+[
+  {
+    "title": "import typing.reveal_type",
+    "kind": "quickfix",
+    "diagnostics": [
+      {
+        "range": {
+          "start": {
+            "line": 0,
+            "character": 0
+          },
+          "end": {
+            "line": 0,
+            "character": 11
+          }
+        },
+        "severity": 2,
+        "code": "undefined-reveal",
+        "codeDescription": {
+          "href": "https://ty.dev/rules#undefined-reveal"
+        },
+        "source": "ty",
+        "message": "`reveal_type` used without importing it",
+        "relatedInformation": []
+      }
+    ],
+    "edit": {
+      "changes": {
+        "file://<temp_dir>/src/foo.py": [
+          {
+            "range": {
+              "start": {
+                "line": 0,
+                "character": 0
+              },
+              "end": {
+                "line": 0,
+                "character": 0
+              }
+            },
+            "newText": "from typing import reveal_type\n"
+          }
+        ]
+      }
+    },
+    "isPreferred": true
+  },
+  {
+    "title": "Ignore 'undefined-reveal' for this line",
+    "kind": "quickfix",
+    "diagnostics": [
+      {
+        "range": {
+          "start": {
+            "line": 0,
+            "character": 0
+          },
+          "end": {
+            "line": 0,
+            "character": 11
+          }
+        },
+        "severity": 2,
+        "code": "undefined-reveal",
+        "codeDescription": {
+          "href": "https://ty.dev/rules#undefined-reveal"
+        },
+        "source": "ty",
+        "message": "`reveal_type` used without importing it",
+        "relatedInformation": []
+      }
+    ],
+    "edit": {
+      "changes": {
+        "file://<temp_dir>/src/foo.py": [
+          {
+            "range": {
+              "start": {
+                "line": 0,
+                "character": 14
+              },
+              "end": {
+                "line": 0,
+                "character": 14
+              }
+            },
+            "newText": "  # ty:ignore[undefined-reveal]"
+          }
+        ]
+      }
+    },
+    "isPreferred": false
+  }
+]