Add code completions to rope_autoimport

pylsp/plugins/rope_autoimport.py

@@ -1,5 +1,6 @@
 # Copyright 2022- Python Language Server Contributors.
 
+import code
 import logging
 from typing import Any, Dict, Generator, List, Optional, Set, Union
 
@@ -21,13 +22,26 @@
 
 _score_pow = 5
 _score_max = 10**_score_pow
-MAX_RESULTS = 1000
+MAX_RESULTS_COMPLETIONS = 1000
+MAX_RESULTS_CODE_ACTIONS = 10
 
 
 @hookimpl
 def pylsp_settings() -> Dict[str, Dict[str, Dict[str, Any]]]:
     # Default rope_completion to disabled
-    return {"plugins": {"rope_autoimport": {"enabled": False, "memory": False}}}
+    return {
+        "plugins": {
+            "rope_autoimport": {
+                "memory": False,
+                "completions": {
+                    "enabled": False,
+                },
+                "code_actions": {
+                    "enabled": False,
+                },
+            }
+        }
+    }
 
 
 # pylint: disable=too-many-return-statements
@@ -122,6 +136,7 @@ def _process_statements(
     word: str,
     autoimport: AutoImport,
     document: Document,
+    feature: str = "completions",
 ) -> Generator[Dict[str, Any], None, None]:
     for suggestion in suggestions:
         insert_line = autoimport.find_insertion_line(document.source) - 1
@@ -134,14 +149,26 @@ def _process_statements(
         if score > _score_max:
             continue
         # TODO make this markdown
-        yield {
-            "label": suggestion.name,
-            "kind": suggestion.itemkind,
-            "sortText": _sort_import(score),
-            "data": {"doc_uri": doc_uri},
-            "detail": _document(suggestion.import_statement),
-            "additionalTextEdits": [edit],
-        }
+        if feature == "completions":
+            yield {
+                "label": suggestion.name,
+                "kind": suggestion.itemkind,
+                "sortText": _sort_import(score),
+                "data": {"doc_uri": doc_uri},
+                "detail": _document(suggestion.import_statement),
+                "additionalTextEdits": [edit],
+            }
+        elif feature == "code_actions":
+            yield {
+                "title": suggestion.import_statement,
+                "kind": "quickfix",
+                "edit": {"changes": {doc_uri: [edit]}},
+                # data is a supported field for codeAction responses
+                # See https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_codeAction
+                "data": {"sortText": _sort_import(score)},
+            }
+        else:
+            raise ValueError(f"Unknown feature: {feature}")
 
 
 def get_names(script: Script) -> Set[str]:
@@ -175,12 +202,12 @@ def pylsp_completions(
     suggestions = list(autoimport.search_full(word, ignored_names=ignored_names))
     results = list(
         sorted(
-            _process_statements(suggestions, document.uri, word, autoimport, document),
+            _process_statements(suggestions, document.uri, word, autoimport, document, "completions"),
             key=lambda statement: statement["sortText"],
         )
     )
-    if len(results) > MAX_RESULTS:
-        results = results[:MAX_RESULTS]
+    if len(results) > MAX_RESULTS_COMPLETIONS:
+        results = results[:MAX_RESULTS_COMPLETIONS]
     return results
 
 
@@ -206,6 +233,58 @@ def _sort_import(score: int) -> str:
     return "[z" + str(score).rjust(_score_pow, "0")
 
 
+@hookimpl
+def pylsp_code_actions(
+    config: Config,
+    workspace: Workspace,
+    document: Document,
+    range: Dict,
+    context: Dict,
+) -> List[Dict]:
+    """
+    Provide code actions through rope.
+
+    Parameters
+    ----------
+    config : pylsp.config.config.Config
+        Current workspace.
+    workspace : pylsp.workspace.Workspace
+        Current workspace.
+    document : pylsp.workspace.Document
+        Document to apply code actions on.
+    range : Dict
+        Range argument given by pylsp. Not used here.
+    context : Dict
+        CodeActionContext given as dict.
+
+    Returns
+    -------
+      List of dicts containing the code actions.
+    """
+    log.debug(f"textDocument/codeAction: {document} {range} {context}")
+    code_actions = []
+    for diagnostic in context.get("diagnostics", []):
+        if not diagnostic.get("message", "").lower().startswith("undefined name"):
+            continue
+        word = diagnostic.get("message", "").split("`")[1]
+        log.debug(f"autoimport: searching for word: {word}")
+        rope_config = config.settings(document_path=document.path).get("rope", {})
+        autoimport = workspace._rope_autoimport(rope_config)
+        suggestions = list(autoimport.search_full(word))
+        log.debug("autoimport: suggestions: %s", suggestions)
+        results = list(
+            sorted(
+                _process_statements(suggestions, document.uri, word, autoimport, document, "code_actions"),
+                key=lambda statement: statement["data"]["sortText"],
+            )
+        )
+        if len(results) > MAX_RESULTS_CODE_ACTIONS:
+            results = results[:MAX_RESULTS_CODE_ACTIONS]
+        code_actions.extend(results)
+
+    return code_actions
+
+
 def _reload_cache(
     config: Config, workspace: Workspace, files: Optional[List[Document]] = None
 ):

pylsp/python_lsp.py

@@ -385,7 +385,7 @@ def watch_parent_process(pid):
     def m_initialized(self, **_kwargs):
         self._hook("pylsp_initialized")
 
-    def code_actions(self, doc_uri, range, context):
+    def code_actions(self, doc_uri: str, range: Dict, context: Dict):
         return flatten(
             self._hook("pylsp_code_actions", doc_uri, range=range, context=context)
         )

test/plugins/test_autoimport.py

@@ -15,6 +15,7 @@
 from pylsp.plugins.rope_autoimport import _get_score, _should_insert, get_names
 from pylsp.plugins.rope_autoimport import (
     pylsp_completions as pylsp_autoimport_completions,
+    pylsp_code_actions as pylsp_autoimport_code_actions,
 )
 from pylsp.plugins.rope_autoimport import pylsp_initialize
 from pylsp.workspace import Workspace
@@ -37,7 +38,7 @@ def autoimport_workspace(tmp_path_factory) -> Workspace:
         uris.from_fs_path(str(tmp_path_factory.mktemp("pylsp"))), Mock()
     )
     workspace._config = Config(workspace.root_uri, {}, 0, {})
-    workspace._config.update({"rope_autoimport": {"memory": True, "enabled": True}})
+    workspace._config.update({"rope_autoimport": {"memory": True, "completions": {"enabled": True}, "code_actions": {"enabled": True}}})
     pylsp_initialize(workspace._config, workspace)
     yield workspace
     workspace.close()
@@ -161,6 +162,30 @@ def test_autoimport_defined_name(config, workspace):
     assert not check_dict({"label": "List"}, completions)
 
 
+def test_autoimport_code_actions(config, autoimport_workspace):
+    source = "os"
+    autoimport_workspace.put_document(DOC_URI, source=source)
+    doc = autoimport_workspace.get_document(DOC_URI)
+    context = {
+        "diagnostics": [
+            {
+                "range": {
+                    "start": {"line": 0, "character": 0},
+                    "end": {"line": 0, "character": 2},
+                },
+                "message": "Undefined name `os`",
+            }
+        ]
+    }
+    actions = pylsp_autoimport_code_actions(
+        config, autoimport_workspace, doc, None, context
+    )
+    autoimport_workspace.rm_document(DOC_URI)
+    assert any(
+        action.get("title") == "import os" for action in actions
+    )
+
+
 class TestShouldInsert:
     def test_dot(self):
         assert not should_insert("""str.""", 4)
@@ -217,7 +242,7 @@ class sfa:
 
 
 @pytest.mark.skipif(IS_WIN, reason="Flaky on Windows")
-def test_autoimport_for_notebook_document(
+def test_autoimport_completions_for_notebook_document(
     client_server_pair,
 ):
     client, server = client_server_pair
@@ -237,13 +262,22 @@ def test_autoimport_for_notebook_document(
 
     server.m_workspace__did_change_configuration(
         settings={
-            "pylsp": {"plugins": {"rope_autoimport": {"enabled": True, "memory": True}}}
+            "pylsp": {
+                "plugins": {
+                    "rope_autoimport": {
+                        "memory": True,
+                        "completions": {
+                            "enabled": True
+                        },
+                    },
+                }
+            }
         }
     )
     rope_autoimport_settings = server.workspace._config.plugin_settings(
         "rope_autoimport"
     )
-    assert rope_autoimport_settings.get("enabled", False) is True
+    assert rope_autoimport_settings.get("completions", {}).get("enabled", False) is True
     assert rope_autoimport_settings.get("memory", False) is True
 
     # 1.