8000 Make autoimport cache generation non-blocking (#499) · python-lsp/python-lsp-server@de87a80 · GitHub
[go: up one dir, main page]
Skip to content

Commit de87a80

Browse files
tkrabeltkrabel
authored
Make autoimport cache generation non-blocking (#499)
1 parent c428381 commit de87a80

File tree

4 files changed

+89
-26
lines changed

4 files changed

+89
-26
lines changed

‎pylsp/_utils.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import pathlib
99
import re
1010
import threading
11+
import time
1112
from typing import List, Optional
1213

1314
import docstring_to_markdown
@@ -55,6 +56,23 @@ def run():
5556
return wrapper
5657

5758

59+
def throttle(seconds=1):
60+
"""Throttles calls to a function evey `seconds` seconds."""
61+
62+
def decorator(func):
63+
@functools.wraps(func)
64+
def wrapper(*args, **kwargs): # pylint: disable=inconsistent-return-statements
65+
if not hasattr(wrapper, "last_call"):
66+
wrapper.last_call = 0
67+
if time.time() - wrapper.last_call >= seconds:
68+
wrapper.last_call = time.time()
69+
return func(*args, **kwargs)
70+
71+
return wrapper
72+
73+
return decorator
74+
75+
5876
def find_parents(root, path, names):
5977
"""Find files matching the given names relative to the given path.
6078

‎pylsp/plugins/_rope_task_handle.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
from __future__ import annotations
22

33
import logging
4+
45
from typing import Callable, ContextManager, List, Optional, Sequence
56

67
from rope.base.taskhandle import BaseJobSet, BaseTaskHandle
78

89
from pylsp.workspace import Workspace
10+
from pylsp._utils import throttle
911

1012
log = logging.getLogger(__name__)
1113
Report = Callable[[str, int], None]
@@ -55,6 +57,7 @@ def increment(self) -> None:
5557
self.count += 1
5658
self._report()
5759

60+
@throttle(0.5)
5861
def _report(self):
5962
percent = int(self.get_percent_done())
6063
message = f"{self.job_name} {self.done}/{self.count}"

‎pylsp/plugins/rope_autoimport.py

Lines changed: 59 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import logging
44
from typing import Any, Dict, Generator, List, Optional, Set, Union
5+
import threading
56

67
import parso
78
from jedi import Script
@@ -25,6 +26,55 @@
2526
MAX_RESULTS_CODE_ACTIONS = 5
2627

2728

29+
class AutoimportCache:
30+
"""Handles the cache creation."""
31+
32+
def __init__(self):
33+
self.thread = None
34+
35+
def reload_cache(
36+
self,
37+
config: Config,
38+
workspace: Workspace,
39+
files: Optional[List[Document]] = None,
40+
single_thread: Optional[bool] = False,
41+
):
42+
if self.is_blocked():
43+
return
44+
45+
memory: bool = config.plugin_settings("rope_autoimport").get("memory", False)
46+
rope_config = config.settings().get("rope", {})
47+
autoimport = workspace._rope_autoimport(rope_config, memory)
48+
resources: Optional[List[Resource]] = (
49+
None
50+
if files is None
51+
else [document._rope_resource(rope_config) for document in files]
52+
)
53+
54+
if single_thread:
55+
self._reload_cache(workspace, autoimport, resources)
56+
else:
57+
# Creating the cache may take 10-20s for a environment with 5k python modules. That's
58+
# why we decided to move cache creation into its own thread.
59+
self.thread = threading.Thread(
60+
target=self._reload_cache, args=(workspace, autoimport, resources)
61+
)
62+
self.thread.start()
63+
64+
def _reload_cache(
65+
self,
66+
workspace: Workspace,
67+
autoimport: AutoImport,
68+
resources: Optional[List[Resource]] = None,
69+
):
70+
task_handle = PylspTaskHandle(workspace)
71+
autoimport.generate_cache(task_handle=task_handle, resources=resources)
72+
autoimport.generate_modules_cache(task_handle=task_handle)
73+
74+
def is_blocked(self):
75+
return self.thread and self.thread.is_alive()
76+
77+
2878
@hookimpl
2979
def pylsp_settings() -> Dict[str, Dict[str, Dict[str, Any]]]:
3080
# Default rope_completion to disabled
@@ -191,7 +241,7 @@ def pylsp_completions(
191241
not config.plugin_settings("rope_autoimport")
192242
.get("completions", {})
193243
.get("enabled", True)
194-
):
244+
) or cache.is_blocked():
195245
return []
196246

197247
line = document.lines[position["line"]]
@@ -283,7 +333,7 @@ def pylsp_code_actions(
283333
not config.plugin_settings("rope_autoimport")
284334
.get("code_actions", {})
285335
.get("enabled", True)
286-
):
336+
) or cache.is_blocked():
287337
return []
288338

289339
log.debug(f"textDocument/codeAction: {document} {range} {context}")
@@ -319,29 +369,13 @@ def pylsp_code_actions(
319369
return code_actions
320370

321371

322-
def _reload_cache(
323-
config: Config, workspace: Workspace, files: Optional[List[Document]] = None
324-
):
325-
memory: bool = config.plugin_settings("rope_autoimport").get("memory", False)
326-
rope_config = config.settings().get("rope", {})
327-
autoimport = workspace._rope_autoimport(rope_config, memory)
328-
task_handle = PylspTaskHandle(workspace)
329-
resources: Optional[List[Resource]] = (
330-
None
331-
if files is None
332-
else [document._rope_resource(rope_config) for document in files]
333-
)
334-
autoimport.generate_cache(task_handle=task_handle, resources=resources)
335-
autoimport.generate_modules_cache(task_handle=task_handle)
336-
337-
338372
@hookimpl
339373
def pylsp_initialize(config: Config, workspace: Workspace):
340374
"""Initialize AutoImport.
341375
342376
Generates the cache for local and global items.
343377
"""
344-
_reload_cache(config, workspace)
378+
cache.reload_cache(config, workspace)
345379

346380

347381
@hookimpl
@@ -350,13 +384,13 @@ def pylsp_document_did_open(config: Config, workspace: Workspace):
350384
351385
Generates the cache for local and global items.
352386
"""
353-
_reload_cache(config, workspace)
387+
cache.reload_cache(config, workspace)
354388

355389

356390
@hookimpl
357391
def pylsp_document_did_save(config: Config, workspace: Workspace, document: Document):
358392
"""Update the names associated with this document."""
359-
_reload_cache(config, workspace, [document])
393+
cache.reload_cache(config, workspace, [document])
360394

361395

362396
@hookimpl
@@ -368,6 +402,9 @@ def pylsp_workspace_configuration_changed(config: Config, workspace: Workspace):
368402
Generates the cache for local and global items.
369403
"""
370404
if config.plugin_settings("rope_autoimport").get("enabled", False):
371-
_reload_cache(config, workspace)
405+
cache.reload_cache(config, workspace)
372406
else:
373407
log.debug("autoimport: Skipping cache reload.")
408+
409+
410+
cache: AutoimportCache = AutoimportCache()

‎test/plugins/test_autoimport.py

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,13 @@
1515
from pylsp.plugins.rope_autoimport import (
1616
_get_score,
1717
_should_insert,
18+
cache,
1819
get_name_or_module,
1920
get_names,
2021
)
2122
from pylsp.plugins.rope_autoimport import (
2223
pylsp_completions as pylsp_autoimport_completions,
2324
)
24-
from pylsp.plugins.rope_autoimport import pylsp_initialize
2525
from pylsp.workspace import Workspace
2626

2727

@@ -57,7 +57,7 @@ def autoimport_workspace(tmp_path_factory) -> Workspace:
5757
}
5858
}
5959
)
60-
pylsp_initialize(workspace._config, workspace)
60+
cache.reload_cache(workspace._config, workspace, single_thread=True)
6161
yield workspace
6262
workspace.close()
6363

@@ -293,7 +293,6 @@ def test_autoimport_code_actions_and_completions_for_notebook_document(
293293
}
294294
},
295295
)
296-
297296
with patch.object(server._endpoint, "notify") as mock_notify:
298297
# Expectations:
299298
# 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(
305304
# 4. We receive an autoimport suggestion for "sys" because it's not already imported.
306305
# 5. If diagnostics doesn't contain "undefined name ...", we send empty quick fix suggestions.
307306
send_notebook_did_open(client, ["os", "import os\nos", "os", "sys"])
308-
wait_for_condition(lambda: mock_notify.call_count >= 3)
307+
wait_for_condition(lambda: mock_notify.call_count >= 4)
308+
# We received diagnostics messages for every cell
309+
assert all(
310+
"textDocument/publishDiagnostics" in c.args
311+
for c in mock_notify.call_args_list
312+
)
309313

310314
rope_autoimport_settings = server.workspace._config.plugin_settings(
311315
"rope_autoimport"
312316
)
313317
assert rope_autoimport_settings.get("completions", {}).get("enabled", False) is True
314318
assert rope_autoimport_settings.get("memory", False) is True
319+
wait_for_condition(lambda: not cache.thread.is_alive())
315320

316321
# 1.
317322
quick_fixes = server.code_actions("cell_1_uri", {}, make_context("os", 0, 0, 2))

0 commit comments

Comments
 (0)
0