followup: add-import action for reveal_type too

Author: Gankra

crates/ty_ide/src/code_action.rs

@@ -3,7 +3,7 @@ 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::types::UNRESOLVED_REFERENCE;
+use ty_python_semantic::types::{UNDEFINED_REVEAL, UNRESOLVED_REFERENCE};
 
 /// A `QuickFix` Code Action
 #[derive(Debug, Clone)]
@@ -23,7 +23,8 @@ pub fn code_actions(
     let Ok(lint_id) = registry.get(diagnostic_id) else {
         return None;
     };
-    if lint_id.name() == UNRESOLVED_REFERENCE.name() {
+    let lint_name = lint_id.name();
+    if lint_name == UNRESOLVED_REFERENCE.name() || lint_name == UNDEFINED_REVEAL.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;

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": "import typing_extensions.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_extensions import reveal_type\n"
+          }
+        ]
+      }
+    },
+    "isPreferred": true
+  }
+]