From 3b703abfbfb6eff8d3d223264252cf35bd6df745 Mon Sep 17 00:00:00 2001
From: Lars Michelsen <lm@checkmk.com>
Date: Sun, 5 Oct 2025 21:21:05 +0200
Subject: [PATCH] rope_autoimport: continue on name-defined diagnostics
 provided by mypy

Previously the autoimport code action did only act on diagnostics which
message contained "undefined name". In my case this did not trigger at
all so that the plugin failed to provide code actions.

This change now additionally looks at name-defined error codes provided
by mypy.
---
 pylsp/plugins/rope_autoimport.py |  7 ++-
 test/plugins/test_autoimport.py  | 89 ++++++++++++++++++++++++++++++++
 2 files changed, 95 insertions(+), 1 deletion(-)

diff --git a/pylsp/plugins/rope_autoimport.py b/pylsp/plugins/rope_autoimport.py
index 8ba951f7..be51d20e 100644
--- a/pylsp/plugins/rope_autoimport.py
+++ b/pylsp/plugins/rope_autoimport.py
@@ -337,7 +337,12 @@ def pylsp_code_actions(
     log.debug(f"textDocument/codeAction: {document} {range} {context}")
     code_actions = []
     for diagnostic in context.get("diagnostics", []):
-        if "undefined name" not in diagnostic.get("message", "").lower():
+        # Provided by flake8 and ruff (F821)
+        matched_message = "undefined name" in diagnostic.get("message", "").lower()
+        # Provided by mypy
+        matched_code = diagnostic.get("code", "") == "name-defined"
+
+        if not matched_message and not matched_code:
             continue
 
         word = get_name_or_module(document, diagnostic)
diff --git a/test/plugins/test_autoimport.py b/test/plugins/test_autoimport.py
index cbe3dde1..454548b9 100644
--- a/test/plugins/test_autoimport.py
+++ b/test/plugins/test_autoimport.py
@@ -346,3 +346,92 @@ def test_autoimport_code_actions_and_completions_for_notebook_document(
     context = {"diagnostics": [{"message": "A random message"}]}
     quick_fixes = server.code_actions("cell_4_uri", {}, context)
     assert len(quick_fixes) == 0
+
+
+def make_mypy_context(module_name, line, character_start, character_end):
+    return {
+        "diagnostics": [
+            {
+                'range': {
+                    'start': {'line': line, 'character': character_start},
+                    'end': {'line': line, 'character': character_end}
+                },
+                'code': 'name-defined'
+            }
+        ]
+    }
+
+
+@pytest.mark.skipif(IS_WIN, reason="Flaky on Windows")
+def test_autoimport_code_actions_and_completions_based_on_mypy_error(
+    client_server_pair,
+) -> None:
+    client, server = client_server_pair
+    send_initialize_request(
+        client,
+        {
+            "pylsp": {
+                "plugins": {
+                    "rope_autoimport": {
+                        "memory": True,
+                        "enabled": True,
+                        "completions": {"enabled": True},
+                    },
+                }
+            }
+        },
+    )
+    with patch.object(server._endpoint, "notify") as mock_notify:
+        # Expectations:
+        # 1. We receive an autoimport suggestion for "os" in the first cell because
+        #    os is imported after that.
+        # 2. We don't receive an autoimport suggestion for "os" in the second cell because it's
+        #    already imported in the second cell.
+        # 3. We don't receive an autoimport suggestion for "os" in the third cell because it's
+        #    already imported in the second cell.
+        # 4. We receive an autoimport suggestion for "sys" because it's not already imported.
+        # 5. If diagnostics doesn't contain "undefined name ...", we send empty quick fix suggestions.
+        send_notebook_did_open(client, ["os", "import os\nos", "os", "sys"])
+        wait_for_condition(lambda: mock_notify.call_count >= 4)
+        # We received diagnostics messages for every cell
+        assert all(
+            "textDocument/publishDiagnostics" in c.args
+            for c in mock_notify.call_args_list
+        )
+
+    rope_autoimport_settings = server.workspace._config.plugin_settings(
+        "rope_autoimport"
+    )
+    assert rope_autoimport_settings.get("completions", {}).get("enabled", False) is True
+    assert rope_autoimport_settings.get("memory", False) is True
+    wait_for_condition(lambda: not cache.is_blocked())
+
+    # 1.
+    quick_fixes = server.code_actions("cell_1_uri", {}, make_mypy_context("os", 0, 0, 2))
+    assert any(s for s in quick_fixes if contains_autoimport_quickfix(s, "os"))
+
+    completions = server.completions("cell_1_uri", position(0, 2)).get("items")
+    assert any(s for s in completions if contains_autoimport_completion(s, "os"))
+
+    # 2.
+    # We don't test code actions here as in this case, there would be no code actions sent bc
+    # there wouldn't be a diagnostics message.
+    completions = server.completions("cell_2_uri", position(1, 2)).get("items")
+    assert not any(s for s in completions if contains_autoimport_completion(s, "os"))
+
+    # 3.
+    # Same as in 2.
+    completions = server.completions("cell_3_uri", position(0, 2)).get("items")
+    assert not any(s for s in completions if contains_autoimport_completion(s, "os"))
+
+    # 4.
+    quick_fixes = server.code_actions("cell_4_uri", {}, make_context("sys", 0, 0, 3))
+    assert any(s for s in quick_fixes if contains_autoimport_quickfix(s, "sys"))
+
+    completions = server.completions("cell_4_uri", position(0, 3)).get("items")
+    assert any(s for s in completions if contains_autoimport_completion(s, "sys"))
+
+    # 5.
+    context = {"diagnostics": [{"code": "some-code"}]}
+    quick_fixes = server.code_actions("cell_4_uri", {}, context)
+    assert len(quick_fixes) == 0