Make autoimport cache generation non-blocking

pylsp/_utils.py

@@ -10,6 +10,8 @@
 import threading
 from typing import List, Optional
 
+import time
+from functools import wraps
 import docstring_to_markdown
 import jedi
 
@@ -55,6 +57,23 @@ def run():
     return wrapper
 
 
+def throttle(seconds=1):
+    """Throttles calls to a function evey `seconds` seconds."""
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):  # pylint: disable=inconsistent-return-statements
+            if not hasattr(wrapper, "last_call"):
+                wrapper.last_call = 0
+            if time.time() - wrapper.last_call >= seconds:
+                wrapper.last_call = time.time()
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 def find_parents(root, path, names):
     """Find files matching the given names relative to the given path.
 

pylsp/plugins/_rope_task_handle.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
 import logging
+
 from typing import Callable, ContextManager, List, Optional, Sequence
 
 from rope.base.taskhandle import BaseJobSet, BaseTaskHandle
 
 from pylsp.workspace import Workspace
+from pylsp._utils import throttle
 
 log = logging.getLogger(__name__)
 Report = Callable[[str, int], None]
@@ -55,6 +57,7 @@ def increment(self) -> None:
         self.count += 1
         self._report()
 
+    @throttle(0.5)
     def _report(self):
         percent = int(self.get_percent_done())
         message = f"{self.job_name} {self.done}/{self.count}"

pylsp/plugins/rope_autoimport.py

@@ -2,6 +2,7 @@
 
 import logging
 from typing import Any, Dict, Generator, List, Optional, Set, Union
+import threading
 
 import parso
 from jedi import Script
@@ -25,6 +26,55 @@
 MAX_RESULTS_CODE_ACTIONS = 5
 
 
+class AutoimportCache:
+    """Handles the cache creation."""
+
+    def __init__(self):
+        self.thread = None
+
+    def reload_cache(
+        self,
+        config: Config,
+        workspace: Workspace,
+        files: Optional[List[Document]] = None,
+        single_thread: Optional[bool] = False,
+    ):
+        if self.is_blocked():
+            return
+
+        memory: bool = config.plugin_settings("rope_autoimport").get("memory", False)
+        rope_config = config.settings().get("rope", {})
+        autoimport = workspace._rope_autoimport(rope_config, memory)
+        resources: Optional[List[Resource]] = (
+            None
+            if files is None
+            else [document._rope_resource(rope_config) for document in files]
+        )
+
+        if single_thread:
+            self._reload_cache(workspace, autoimport, resources)
+        else:
+            # Creating the cache may take 10-20s for a environment with 5k python modules. That's
+            # why we decided to move cache creation into its own thread.
+            self.thread = threading.Thread(
+                target=self._reload_cache, args=(workspace, autoimport, resources)
+            )
+            self.thread.start()
+
+    def _reload_cache(
+        self,
+        workspace: Workspace,
+        autoimport: AutoImport,
+        resources: Optional[List[Resource]] = None,
+    ):
+        task_handle = PylspTaskHandle(workspace)
+        autoimport.generate_cache(task_handle=task_handle, resources=resources)
+        autoimport.generate_modules_cache(task_handle=task_handle)
+
+    def is_blocked(self):
+        return self.thread and self.thread.is_alive()
+
+
 @hookimpl
 def pylsp_settings() -> Dict[str, Dict[str, Dict[str, Any]]]:
     # Default rope_completion to disabled
@@ -191,7 +241,7 @@ def pylsp_completions(
         not config.plugin_settings("rope_autoimport")
         .get("completions", {})
         .get("enabled", True)
-    ):
+    ) or cache.is_blocked():
         return []
 
     line = document.lines[position["line"]]
@@ -283,7 +333,7 @@ def pylsp_code_actions(
         not config.plugin_settings("rope_autoimport")
         .get("code_actions", {})
         .get("enabled", True)
-    ):
+    ) or cache.is_blocked():
         return []
 
     log.debug(f"textDocument/codeAction: {document} {range} {context}")
@@ -319,29 +369,13 @@ def pylsp_code_actions(
     return code_actions
 
 
-def _reload_cache(
-    config: Config, workspace: Workspace, files: Optional[List[Document]] = None
-):
-    memory: bool = config.plugin_settings("rope_autoimport").get("memory", False)
-    rope_config = config.settings().get("rope", {})
-    autoimport = workspace._rope_autoimport(rope_config, memory)
-    task_handle = PylspTaskHandle(workspace)
-    resources: Optional[List[Resource]] = (
-        None
-        if files is None
-        else [document._rope_resource(rope_config) for document in files]
-    )
-    autoimport.generate_cache(task_handle=task_handle, resources=resources)
-    autoimport.generate_modules_cache(task_handle=task_handle)
-
-
 @hookimpl
 def pylsp_initialize(config: Config, workspace: Workspace):
     """Initialize AutoImport.
 
     Generates the cache for local and global items.
     """
-    _reload_cache(config, workspace)
+    cache.reload_cache(config, workspace)
 
 
 @hookimpl
@@ -350,13 +384,13 @@ def pylsp_document_did_open(config: Config, workspace: Workspace):
 
     Generates the cache for local and global items.
     """
-    _reload_cache(config, workspace)
+    cache.reload_cache(config, workspace)
 
 
 @hookimpl
 def pylsp_document_did_save(config: Config, workspace: Workspace, document: Document):
     """Update the names associated with this document."""
-    _reload_cache(config, workspace, [document])
+    cache.reload_cache(config, workspace, [document])
 
 
 @hookimpl
@@ -368,6 +402,9 @@ def pylsp_workspace_configuration_changed(config: Config, workspace: Workspace):
     Generates the cache for local and global items.
     """
     if config.plugin_settings("rope_autoimport").get("enabled", False):
-        _reload_cache(config, workspace)
+        cache.reload_cache(config, workspace)
     else:
         log.debug("autoimport: Skipping cache reload.")
+
+
+cache: AutoimportCache = AutoimportCache()

test/plugins/test_autoimport.py

@@ -15,13 +15,13 @@
 from pylsp.plugins.rope_autoimport import (
     _get_score,
     _should_insert,
+    cache,
     get_name_or_module,
     get_names,
 )
 from pylsp.plugins.rope_autoimport import (
     pylsp_completions as pylsp_autoimport_completions,
 )
-from pylsp.plugins.rope_autoimport import pylsp_initialize
 from pylsp.workspace import Workspace
 
 
@@ -57,7 +57,7 @@ def autoimport_workspace(tmp_path_factory) -> Workspace:
             }
         }
     )
-    pylsp_initialize(workspace._config, workspace)
+    cache.reload_cache(workspace._config, workspace, single_thread=True)
     yield workspace
     workspace.close()
 
@@ -293,7 +293,6 @@ def test_autoimport_code_actions_and_completions_for_notebook_document(
             }
         },
     )
-
     with patch.object(server._endpoint, "notify") as mock_notify:
         # Expectations:
         # 1. We receive an autoimport suggestion for "os" in the first cell because
@@ -305,13 +304,19 @@ def test_autoimport_code_actions_and_completions_for_notebook_document(
         # 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 >= 3)
+        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.thread.is_alive())
 
     # 1.
     quick_fixes = server.code_actions("cell_1_uri", {}, make_context("os", 0, 0, 2))