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.

Author: LarsMichelsen

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)

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