Add a plugin to provide autoimport functionality

CONFIGURATION.md

@@ -59,6 +59,7 @@ This server can be configured using `workspace/didChangeConfiguration` method. E
 | `pylsp.plugins.pylint.args` | `array`  of non-unique `string` items | Arguments to pass to pylint. | `null` |
 | `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_completion.enabled` | `boolean` | Enable or disable the plugin. | `false` |
+| `pylsp.plugins.rope_autoimport.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` |
 | `pylsp.rope.extensionModules` | `string` | Builtin and c-extension modules that are allowed to be imported and inspected by rope. | `null` |

README.md

@@ -95,6 +95,7 @@ As an example, to change the list of errors that pycodestyle will ignore, assumi
 ## LSP Server Features
 
 * Auto Completion
+* [Autoimport](docs/autoimport.md)
 * Code Linting
 * Signature Help
 * Go to definition

docs/autoimport.md

@@ -0,0 +1,18 @@
+# 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/schema.json

@@ -335,6 +335,10 @@
       "type": "boolean",
       "default": false,
       "description": "Enable or disable the plugin."
+    },"pylsp.plugins.rope_autoimport.enabled": {
+      "type": "boolean",
+      "default": false,
+      "description": "Enable or disable the plugin."
     },
     "pylsp.plugins.rope_completion.eager": {
       "type": "boolean",

pylsp/plugins/rope_autoimport.py

@@ -0,0 +1,234 @@
+import logging
+from functools import lru_cache
+from typing import Any, Dict, Generator, List, Set
+
+import parso
+from parso.python import tree
+from parso.tree import NodeOrLeaf
+from rope.base.resources import Resource
+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
+
+
+@hookimpl
+def pylsp_settings() -> Dict[str, Dict[str, Dict[str, Any]]]:
+    # Default rope_completion to disabled
+    return {"plugins": {"rope_autoimport": {"enabled": 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 == "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,
+    doc_uri: str,
+    word: str,
+    autoimport: AutoImport,
+    document: Document,
+) -> Generator[Dict[str, Any], None, None]:
+    for import_statement, name, source, itemkind 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": import_statement + "\n"}
+        score = _get_score(source, import_statement, name, word)
+        if score > _score_max:
+            continue
+        yield {
+            "label": name,
+            "kind": itemkind,
+            "sortText": _sort_import(score),
+            "data": {"doc_uri": doc_uri},
+            "detail": _document(import_statement),
+            "additionalTextEdits": [edit],
+        }
+
+
+def _get_names_from_import(node: tree.Import) -> Generator[str, None, None]:
+    if not node.is_star_import():
+        for name in node.children:
+            if isinstance(name, tree.PythonNode):
+                for sub_name in name.children:
+                    if isinstance(sub_name, tree.Name):
+                        yield sub_name.value
+            elif isinstance(name, tree.Name):
+                yield name.value
+
+
+@lru_cache(maxsize=100)
+def get_names(file: str) -> Generator[str, None, None]:
+    """Get all names to ignore from the current file."""
+    expr = parso.parse(file)
+    for item in expr.children:
+        if isinstance(item, tree.PythonNode):
+            for child in item.children:
+                if isinstance(child, (tree.ImportFrom, tree.ExprStmt)):
+                    for name in child.get_defined_names():
+                        yield name.value
+                elif isinstance(child, tree.Import):
+                    for name in _get_names_from_import(child):
+                        yield name
+
+        if isinstance(item, (tree.Function, tree.Class)):
+            yield item.name.value
+
+
+@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
+    rope_config = config.settings(document_path=document.path).get("rope", {})
+    rope_project = workspace._rope_project_builder(rope_config)
+    ignored_names: Set[str] = set(get_names(document.source))
+    autoimport = AutoImport(rope_project, memory=False)
+    suggestions = list(autoimport.search_full(word, ignored_names=ignored_names))
+    autoimport.close()
+    results = list(
+        sorted(
+            _process_statements(suggestions, document.uri, word, autoimport, document),
+            key=lambda statement: statement["sortText"],
+        )
+    )
+    max_size = 100
+    if len(results) > max_size:
+        results = results[:max_size]
+    return results
+
+
+def _document(import_statement: str) -> str:
+    return "__autoimport__\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."""
+    rope_config = config.settings().get("rope", {})
+    rope_project = workspace._rope_project_builder(rope_config)
+    autoimport = AutoImport(rope_project, memory=False)
+    autoimport.generate_modules_cache()
+    autoimport.generate_cache()
+    autoimport.close()
+
+
+@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)
+    rope_project = workspace._rope_project_builder(rope_config)
+    autoimport = AutoImport(rope_project, memory=False)
+    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()
+    autoimport.close()

pylsp/python_lsp.py

@@ -274,6 +274,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)
 
@@ -340,6 +343,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

@@ -56,7 +56,10 @@ 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)
+            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)

pyproject.toml

@@ -33,7 +33,7 @@ all = [
     "pydocstyle>=2.0.0",
     "pyflakes>=2.4.0,<2.5.0",
     "pylint>=2.5.0",
-    "rope>=0.10.5",
+    "rope>=1.1.1",
     "yapf",
 ]
 autopep8 = ["autopep8>=1.6.0,<1.7.0"]
@@ -43,7 +43,7 @@ pycodestyle = ["pycodestyle>=2.8.0,<2.9.0"]
 pydocstyle = ["pydocstyle>=2.0.0"]
 pyflakes = ["pyflakes>=2.4.0,<2.5.0"]
 pylint = ["pylint>=2.5.0"]
-rope = ["rope>0.10.5"]
+rope = ["rope>1.1.1"]
 yapf = ["yapf"]
 test = [
     "pylint>=2.5.0",
@@ -77,6 +77,7 @@ pyflakes = "pylsp.plugins.pyflakes_lint"
 pylint = "pylsp.plugins.pylint_lint"
 rope_completion = "pylsp.plugins.rope_completion"
 rope_rename = "pylsp.plugins.rope_rename"
+rope_autoimport = "pylsp.plugins.rope_autoimport"
 yapf = "pylsp.plugins.yapf_format"
 
 [project.scripts]