diff --git a/pylsp/_utils.py b/pylsp/_utils.py index 9d393b92..f694a36a 100644 --- a/pylsp/_utils.py +++ b/pylsp/_utils.py @@ -8,6 +8,7 @@ import pathlib import re import threading +import time from typing import List, Optional import docstring_to_markdown @@ -55,6 +56,23 @@ def run(): return wrapper +def throttle(seconds=1): + """Throttles calls to a function evey `seconds` seconds.""" + + def decorator(func): + @functools.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. diff --git a/pylsp/plugins/_rope_task_handle.py b/pylsp/plugins/_rope_task_handle.py index ed93a058..841d6fee 100644 --- a/pylsp/plugins/_rope_task_handle.py +++ b/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}" diff --git a/pylsp/plugins/rope_autoimport.py b/pylsp/plugins/rope_autoimport.py index ebcc7070..b3626408 100644 --- a/pylsp/plugins/rope_autoimport.py +++ b/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() diff --git a/test/plugins/test_autoimport.py b/test/plugins/test_autoimport.py index 9f02965b..6d3eeb50 100644 --- a/test/plugins/test_autoimport.py +++ b/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))