Add a plugin to provide autoimport functionality

CONFIGURATION.md

@@ -60,6 +60,8 @@ This server can be configured using `workspace/didChangeConfiguration` method. E
 | `pylsp.plugins.pylint.enabled` | `boolean` | Enable or disable the plugin. | `false` |
 | `pylsp.plugins.pylint.args` | `array` of non-unique `string` items | Arguments to pass to pylint. | `[]` |
 | `pylsp.plugins.pylint.executable` | `string` | Executable to run pylint with. Enabling this will run pylint on unsaved files via stdin. Can slow down workflow. Only works with python3. | `null` |
+| `pylsp.plugins.rope_autoimport.enabled` | `boolean` | Enable or disable autoimport. | `false` |
+| `pylsp.plugins.rope_autoimport.memory` | `boolean` | Make the autoimport database memory only. Drastically increases startup time. | `false` |
 | `pylsp.plugins.rope_completion.enabled` | `boolean` | Enable or disable the plugin. | `false` |
 | `pylsp.plugins.rope_completion.eager` | `boolean` | Resolve documentation and detail eagerly. | `false` |
 | `pylsp.plugins.yapf.enabled` | `boolean` | Enable or disable the plugin. | `true` |

README.md

@@ -140,6 +140,7 @@ pip install 'python-lsp-server[websockets]'
 ## LSP Server Features
 
 * Auto Completion
+* [Autoimport](docs/autoimport.md)
 * Code Linting
 * Signature Help
 * Go to definition

docs/autoimport.md

@@ -0,0 +1,25 @@
+# Autoimport for pylsp
+
+Requirements:
+
+1. install `python-lsp-server[rope]`
+2. set `pylsp.plugins.rope_autoimport.enabled` to `true`
+
+## Startup
+
+Autoimport will generate an autoimport sqllite3 database in .ropefolder/autoimport.db on startup.  
+This will take a few seconds but should be much quicker on future runs.
+
+## Usage
+
+Autoimport will provide suggestions to import names from everything in `sys.path`. You can change this by changing where pylsp is running or by setting rope's 'python_path' option.
+It will suggest modules, submodules, keywords, functions, and classes.
+
+Since autoimport inserts everything towards the end of the import group, its recommended you use the isort [plugin](https://github.com/paradoxxxzero/pyls-isort).
+
+## Credits
+
+- Most of the code was written by me, @bageljrkhanofemus
+- [lyz-code](https://github.com/lyz-code/autoimport) for inspiration and some ideas
+- [rope](https://github.com/python-rope/rope), especially @lieryan
+- [pyright](https://github.com/Microsoft/pyright) for details on language server implementation

pylsp/config/config.py

@@ -81,7 +81,6 @@ def __init__(self, root_uri, init_opts, process_id, capabilities):
             if plugin is not None:
                 log.info("Loaded pylsp plugin %s from %s", name, plugin)
 
-        # pylint: disable=no-member
         for plugin_conf in self._pm.hook.pylsp_settings(config=self):
             self._plugin_settings = _utils.merge_dicts(self._plugin_settings, plugin_conf)
 

pylsp/config/schema.json

@@ -354,6 +354,16 @@
       "default": null,
       "description": "Executable to run pylint with. Enabling this will run pylint on unsaved files via stdin. Can slow down workflow. Only works with python3."
     },
+    "pylsp.plugins.rope_autoimport.enabled": {
+      "type": "boolean",
+      "default": false,
+      "description": "Enable or disable autoimport."
+    },
+    "pylsp.plugins.rope_autoimport.memory": {
+      "type": "boolean",
+      "default": false,
+      "description": "Make the autoimport database memory only. Drastically increases startup time."
+    },
     "pylsp.plugins.rope_completion.enabled": {
       "type": "boolean",
       "default": false,

pylsp/plugins/rope_autoimport.py

@@ -0,0 +1,219 @@
+# Copyright 2022- Python Language Server Contributors.
+
+import logging
+from typing import Any, Dict, Generator, List, Set
+
+import parso
+from jedi import Script
+from parso.python import tree
+from parso.tree import NodeOrLeaf
+from rope.base.resources import Resource
+from rope.contrib.autoimport.defs import SearchResult
+from rope.contrib.autoimport.sqlite import AutoImport
+
+from pylsp import hookimpl
+from pylsp.config.config import Config
+from pylsp.workspace import Document, Workspace
+
+log = logging.getLogger(__name__)
+
+_score_pow = 5
+_score_max = 10**_score_pow
+MAX_RESULTS = 1000
+
+
+@hookimpl
+def pylsp_settings() -> Dict[str, Dict[str, Dict[str, Any]]]:
+    # Default rope_completion to disabled
+    return {"plugins": {"rope_autoimport": {"enabled": True, "memory": False}}}
+
+
+def _should_insert(expr: tree.BaseNode, word_node: tree.Leaf) -> bool:
+    """
+    Check if we should insert the word_node on the given expr.
+
+    Works for both correct and incorrect code. This is because the
+    user is often working on the code as they write it.
+    """
+    if len(expr.children) == 0:
+        return True
+    first_child = expr.children[0]
+    if isinstance(first_child, tree.EndMarker):
+        if "#" in first_child.prefix:
+            return False  # Check for single line comment
+    if first_child == word_node:
+        return True  # If the word is the first word then its fine
+    if len(expr.children) > 1:
+        if any(node.type == "operator" and "." in node.value or
+               node.type == "trailer" for node in expr.children):
+            return False  # Check if we're on a method of a function
+    if isinstance(first_child, (tree.PythonErrorNode, tree.PythonNode)):
+        # The tree will often include error nodes like this to indicate errors
+        # we want to ignore errors since the code is being written
+        return _should_insert(first_child, word_node)
+    return _handle_first_child(first_child, expr, word_node)
+
+
+def _handle_first_child(first_child: NodeOrLeaf, expr: tree.BaseNode,
+                        word_node: tree.Leaf) -> bool:
+    """Check if we suggest imports given the following first child."""
+    if isinstance(first_child, tree.Import):
+        return False
+    if isinstance(first_child, (tree.PythonLeaf, tree.PythonErrorLeaf)):
+        # Check if the first item is a from or import statement even when incomplete
+        if first_child.value in ("import", "from"):
+            return False
+    if isinstance(first_child, tree.Keyword):
+        if first_child.value == "def":
+            return _should_import_function(word_node, expr)
+        if first_child.value == "class":
+            return _should_import_class(word_node, expr)
+    return True
+
+
+def _should_import_class(word_node: tree.Leaf, expr: tree.BaseNode) -> bool:
+    prev_node = None
+    for node in expr.children:
+        if isinstance(node, tree.Name):
+            if isinstance(prev_node, tree.Operator):
+                if node == word_node and prev_node.value == "(":
+                    return True
+        prev_node = node
+
+    return False
+
+
+def _should_import_function(word_node: tree.Leaf, expr: tree.BaseNode) -> bool:
+    prev_node = None
+    for node in expr.children:
+        if _handle_argument(node, word_node):
+            return True
+        if isinstance(prev_node, tree.Operator):
+            if prev_node.value == "->":
+                if node == word_node:
+                    return True
+        prev_node = node
+    return False
+
+
+def _handle_argument(node: NodeOrLeaf, word_node: tree.Leaf):
+    if isinstance(node, tree.PythonNode):
+        if node.type == "tfpdef":
+            if node.children[2] == word_node:
+                return True
+        if node.type == "parameters":
+            for parameter in node.children:
+                if _handle_argument(parameter, word_node):
+                    return True
+    return False
+
+
+def _process_statements(
+    suggestions: List[SearchResult],
+    doc_uri: str,
+    word: str,
+    autoimport: AutoImport,
+    document: Document,
+) -> Generator[Dict[str, Any], None, None]:
+    for suggestion in suggestions:
+        insert_line = autoimport.find_insertion_line(document.source) - 1
+        start = {"line": insert_line, "character": 0}
+        edit_range = {"start": start, "end": start}
+        edit = {
+            "range": edit_range,
+            "newText": suggestion.import_statement + "\n"
+        }
+        score = _get_score(suggestion.source, suggestion.import_statement,
+                           suggestion.name, word)
+        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],
+        }
+
+
+def get_names(script: Script) -> Set[str]:
+    """Get all names to ignore from the current file."""
+    raw_names = script.get_names(definitions=True)
+    log.debug(raw_names)
+    return set(name.name for name in raw_names)
+
+
+@hookimpl
+def pylsp_completions(config: Config, workspace: Workspace, document: Document,
+                      position):
+    """Get autoimport suggestions."""
+    line = document.lines[position["line"]]
+    expr = parso.parse(line)
+    word_node = expr.get_leaf_for_position((1, position["character"]))
+    if not _should_insert(expr, word_node):
+        return []
+    word = word_node.value
+    log.debug(f"autoimport: searching for word: {word}")
+    rope_config = config.settings(document_path=document.path).get("rope", {})
+    ignored_names: Set[str] = get_names(
+        document.jedi_script(use_document_path=True))
+    autoimport = workspace._rope_autoimport(rope_config)
+    suggestions = list(
+        autoimport.search_full(word, ignored_names=ignored_names))
+    results = list(
+        sorted(
+            _process_statements(suggestions, document.uri, word, autoimport,
+                                document),
+            key=lambda statement: statement["sortText"],
+        ))
+    if len(results) > MAX_RESULTS:
+        results = results[:MAX_RESULTS]
+    return results
+
+
+def _document(import_statement: str) -> str:
+    return """# Auto-Import\n""" + import_statement
+
+
+def _get_score(source: int, full_statement: str, suggested_name: str,
+               desired_name) -> int:
+    import_length = len("import")
+    full_statement_score = len(full_statement) - import_length
+    suggested_name_score = ((len(suggested_name) - len(desired_name)))**2
+    source_score = 20 * source
+    return suggested_name_score + full_statement_score + source_score
+
+
+def _sort_import(score: int) -> str:
+    score = max(min(score, (_score_max) - 1), 0)
+    # Since we are using ints, we need to pad them.
+    # We also want to prioritize autoimport behind everything since its the last priority.
+    # The minimum is to prevent score from overflowing the pad
+    return "[z" + str(score).rjust(_score_pow, "0")
+
+
+@hookimpl
+def pylsp_initialize(config: Config, workspace: Workspace):
+    """Initialize AutoImport. Generates the cache for local and global items."""
+    memory: bool = config.plugin_settings("rope_autoimport").get(
+        "memory", False)
+    rope_config = config.settings().get("rope", {})
+    autoimport = workspace._rope_autoimport(rope_config, memory)
+    autoimport.generate_modules_cache()
+    autoimport.generate_cache()
+
+
+@hookimpl
+def pylsp_document_did_save(config: Config, workspace: Workspace,
+                            document: Document):
+    """Update the names associated with this document."""
+    rope_config = config.settings().get("rope", {})
+    rope_doucment: Resource = document._rope_resource(rope_config)
+    autoimport = workspace._rope_autoimport(rope_config)
+    autoimport.generate_cache(resources=[rope_doucment])
+    # Might as well using saving the document as an indicator to regenerate the module cache
+    autoimport.generate_modules_cache()

pylsp/python_lsp.py

@@ -34,7 +34,6 @@ class _StreamHandlerWrapper(socketserver.StreamRequestHandler):
 
     def setup(self):
         super().setup()
-        # pylint: disable=no-member
         self.delegate = self.DELEGATE_CLASS(self.rfile, self.wfile)
 
     def handle(self):
@@ -48,7 +47,6 @@ def handle(self):
                 if isinstance(e, WindowsError) and e.winerror == 10054:
                     pass
 
-        # pylint: disable=no-member
         self.SHUTDOWN_CALL()
 
 
@@ -212,6 +210,8 @@ def __getitem__(self, item):
         raise KeyError()
 
     def m_shutdown(self, **_kwargs):
+        for workspace in self.workspaces.values():
+            workspace.close()
         self._shutdown = True
 
     def m_exit(self, **_kwargs):
@@ -351,6 +351,9 @@ def definitions(self, doc_uri, position):
     def document_symbols(self, doc_uri):
         return flatten(self._hook('pylsp_document_symbols', doc_uri))
 
+    def document_did_save(self, doc_uri):
+        return self._hook("pylsp_document_did_save", doc_uri)
+
     def execute_command(self, command, arguments):
         return self._hook('pylsp_execute_command', command=command, arguments=arguments)
 
@@ -417,6 +420,7 @@ def m_text_document__did_change(self, contentChanges=None, textDocument=None, **
 
     def m_text_document__did_save(self, textDocument=None, **_kwargs):
         self.lint(textDocument['uri'], is_saved=True)
+        self.document_did_save(textDocument['uri'])
 
     def m_text_document__code_action(self, textDocument=None, range=None, context=None, **_kwargs):
         return self.code_actions(textDocument['uri'], range, context)

pylsp/workspace.py

@@ -7,6 +7,7 @@
 import re
 import functools
 from threading import RLock
+from typing import Optional
 
 import jedi
 
@@ -48,6 +49,15 @@ def __init__(self, root_uri, endpoint, config=None):
         # Whilst incubating, keep rope private
         self.__rope = None
         self.__rope_config = None
+        self.__rope_autoimport = None
+
+    def _rope_autoimport(self, rope_config: Optional, memory: bool = False):
+        # pylint: disable=import-outside-toplevel
+        from rope.contrib.autoimport.sqlite import AutoImport
+        if self.__rope_autoimport is None:
+            project = self._rope_project_builder(rope_config)
+            self.__rope_autoimport = AutoImport(project, memory=memory)
+        return self.__rope_autoimport
 
     def _rope_project_builder(self, rope_config):
         # pylint: disable=import-outside-toplevel
@@ -56,8 +66,12 @@ def _rope_project_builder(self, rope_config):
         # TODO: we could keep track of dirty files and validate only those
         if self.__rope is None or self.__rope_config != rope_config:
             rope_folder = rope_config.get('ropeFolder')
-            self.__rope = Project(self._root_path, ropefolder=rope_folder)
-            self.__rope.prefs.set('extension_modules', rope_config.get('extensionModules', []))
+            if rope_folder:
+                self.__rope = Project(self._root_path, ropefolder=rope_folder)
+            else:
+                self.__rope = Project(self._root_path)
+            self.__rope.prefs.set('extension_modules',
+                                  rope_config.get('extensionModules', []))
             self.__rope.prefs.set('ignore_syntax_errors', True)
             self.__rope.prefs.set('ignore_bad_imports', True)
         self.__rope.validate()
@@ -128,6 +142,10 @@ def _create_document(self, doc_uri, source=None, version=None):
             rope_project_builder=self._rope_project_builder,
         )
 
+    def close(self):
+        if self.__rope_autoimport is not None:
+            self.__rope_autoimport.close()
+
 
 class Document: