Report autoimport progress

pylsp/plugins/_rope_task_handle.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import logging
+from typing import List, Optional, Sequence, Callable, ContextManager
+from rope.base.taskhandle import BaseJobSet, BaseTaskHandle
+
+from pylsp.workspace import Workspace
+
+
+log = logging.getLogger(__name__)
+Report = Callable[[str, int], None]
+
+
+class PylspJobSet(BaseJobSet):
+    count: int = 0
+    done: int = 0
+    _reporter: Report
+    _report_iter: ContextManager
+    job_name: str = ""
+
+    def __init__(self, count: Optional[int], report_iter: ContextManager):
+        if count is not None:
+            self.count = count
+        self._reporter = report_iter.__enter__()
+        self._report_iter = report_iter
+
+    def started_job(self, name: Optional[str]) -> None:
+        if name:
+            self.job_name = name
+
+    def finished_job(self) -> None:
+        self.done += 1
+        if self.get_percent_done() is not None and int(self.get_percent_done()) >= 100:
+            if self._report_iter is None:
+                return
+            self._report_iter.__exit__(None, None, None)
+            self._report_iter = None
+        else:
+            self._report()
+
+    def check_status(self) -> None:
+        pass
+
+    def get_percent_done(self) -> Optional[float]:
+        if self.count == 0:
+            return 0
+        return (self.done / self.count) * 100
+
+    def increment(self) -> None:
+        """
+        Increment the number of tasks to complete.
+        This is used if the number is not known ahead of time.
+        """
+        self.count += 1
+        self._report()
+
+    def _report(self):
+        percent = int(self.get_percent_done())
+        message = f"{self.job_name} {self.done}/{self.count}"
+        log.debug(f"Reporting {message} {percent}%")
+        self._reporter(message, percent)
+
+
+class PylspTaskHandle(BaseTaskHandle):
+    name: str
+    observers: List
+    job_sets: List[PylspJobSet]
+    stopped: bool
+    workspace: Workspace
+    _report: Callable[[str, str], None]
+
+    def __init__(self, workspace: Workspace):
+        self.workspace = workspace
+        self.job_sets = []
+        self.observers = []
+
+    def create_jobset(self, name="JobSet", count: Optional[int] = None):
+        report_iter = self.workspace.report_progress(name, None, None)
+        result = PylspJobSet(count, report_iter)
+        self.job_sets.append(result)
+        self._inform_observers()
+        return result
+
+    def stop(self) -> None:
+        pass
+
+    def current_jobset(self) -> Optional[BaseJobSet]:
+        pass
+
+    def add_observer(self) -> None:
+        pass
+
+    def is_stopped(self) -> bool:
+        pass
+
+    def get_jobsets(self) -> Sequence[BaseJobSet]:
+        pass
+
+    def _inform_observers(self) -> None:
+        for observer in self.observers:
+            observer()

pylsp/plugins/rope_autoimport.py

@@ -1,7 +1,7 @@
 # Copyright 2022- Python Language Server Contributors.
 
 import logging
-from typing import Any, Dict, Generator, List, Set
+from typing import Any, Dict, Generator, List, Set, Optional
 
 import parso
 from jedi import Script
@@ -14,7 +14,7 @@
 from pylsp import hookimpl
 from pylsp.config.config import Config
 from pylsp.workspace import Document, Workspace
-
+from ._rope_task_handle import PylspTaskHandle
 log = logging.getLogger(__name__)
 
 _score_pow = 5
@@ -196,24 +196,32 @@ def _sort_import(score: int) -> str:
     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."""
+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)
-    autoimport.generate_modules_cache()
-    autoimport.generate_cache()
+    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)
+
+
+@hookimpl
+def pylsp_document_did_open(config: Config, workspace: Workspace):
+    """Initialize AutoImport. Generates the cache for local and global items."""
+    _reload_cache(config, workspace)
 
 
 @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()
+    _reload_cache(config, workspace, [document])